Update lunacode.py

lunacode.py

@@ -5,10 +5,31 @@ from duckduckgo_search import DDGS
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import wikipedia
 
-# ✅ Load Luna model from Hugging Face
 model_path = "cosmosai471/Luna-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).to("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModelForCausalLM.from_pretrained(
+    model_path, torch_dtype=torch.bfloat16
+).to("cuda" if torch.cuda.is_available() else "cpu")
+
+def detect_mode(query):
+    code_keywords = ["code", "program", "python", "javascript", "function", "script", "build", "html", "css"]
+    creative_keywords = ["story", "write a story", "poem", "creative", "imagine", "novel", "dialogue"]
+
+    if any(kw in query.lower() for kw in code_keywords):
+        return "code"
+    elif any(kw in query.lower() for kw in creative_keywords):
+        return "creative"
+    else:
+        return "general"
+
+def get_generation_params(query):
+    mode = detect_mode(query)
+    if mode == "code":
+        return 0.3, 0.85
+    elif mode == "creative":
+        return 0.95, 0.95
+    else:
+        return 0.7, 0.9
 
 def luna_generate(prompt, max_tokens=200, temperature=0.7, top_p=0.95):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -21,10 +42,6 @@ def luna_generate(prompt, max_tokens=200, temperature=0.7, top_p=0.95):
     )
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
-def is_code_related(question):
-    keywords = ["code", "program", "python", "javascript", "build a function", "algorithm", "write a", "html", "css"]
-    return any(kw in question.lower() for kw in keywords)
-
 def code_prompt_from_question(question):
     return f'''You are a helpful AI programmer. Your task is to generate complete and clean code with explanations.
 
@@ -72,9 +89,10 @@ def scrape_first_good_content(urls):
             continue
     return None, None
 
-def smart_luna_answer(user_question, max_tokens=512, temperature=0.7, top_p=0.95):
-    # Code check
-    if is_code_related(user_question):
+def smart_luna_answer(user_question, max_tokens=512):
+    temperature, top_p = get_generation_params(user_question)
+
+    if detect_mode(user_question) == "code":
         prompt = code_prompt_from_question(user_question)
         code_response = luna_generate(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
 
@@ -96,7 +114,6 @@ Code:
                 return f"Luna (code + web-enhanced):\n{code_response.strip()}\n\n(Source: {url_used})"
         return f"Luna (code):\n{code_response.strip()}"
 
-    # Base Luna response
     base_prompt = f"User: {user_question}\nLuna:"
     base_answer = luna_generate(base_prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)