Update lunacode.py

lunacode.py

@@ -5,23 +5,26 @@ from duckduckgo_search import DDGS
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import wikipedia
 
-# ✅ Load Luna model from Google Drive
-model_path = "/content/drive/MyDrive/Luna_1"
+# ✅ 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")
-
-# ✅ Generate response from Luna
-def luna_generate(prompt, max_tokens=200):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    output = model.generate(**inputs, max_new_tokens=max_tokens, do_sample=True, temperature=0.7)
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).to("cuda" if torch.cuda.is_available() else "cpu")
+
+def luna_generate(prompt, max_tokens=200, temperature=0.7, top_p=0.95):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    output = model.generate(
+        **inputs,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p
+    )
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
-# ✅ Check if query is about coding
 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)
 
-# ✅ Format prompt for coding
 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.
 
@@ -36,7 +39,6 @@ Requirements:
 Code:
 '''
 
-# ✅ Wikipedia fallback first
 def try_wikipedia(query):
     try:
         wikipedia.set_lang("en")
@@ -50,7 +52,6 @@ def try_wikipedia(query):
         except:
             return None, None
 
-# ✅ DuckDuckGo search
 def get_top_webpages(query, max_results=3):
     urls = []
     with DDGS() as ddgs:
@@ -59,7 +60,6 @@ def get_top_webpages(query, max_results=3):
             urls.append(result["href"])
     return urls
 
-# ✅ Web scraper
 def scrape_first_good_content(urls):
     for url in urls:
         try:
@@ -72,14 +72,12 @@ def scrape_first_good_content(urls):
             continue
     return None, None
 
-# ✅ Smart Luna answer with code + current info
-def smart_luna_answer(user_question):
-    # Step 0: Check if it's code related
+def smart_luna_answer(user_question, max_tokens=512, temperature=0.7, top_p=0.95):
+    # Code check
     if is_code_related(user_question):
         prompt = code_prompt_from_question(user_question)
-        code_response = luna_generate(prompt, max_tokens=512)
+        code_response = luna_generate(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
 
-        # If code is too short or seems weak, enhance with web context
         if len(code_response.strip()) < 100 or "not sure" in code_response.lower():
             urls = get_top_webpages(user_question)
             context, url_used = scrape_first_good_content(urls)
@@ -94,17 +92,15 @@ Generate well-structured code with inline comments and clear function design. On
 
 Code:
 '''
-                code_response = luna_generate(enhanced_prompt, max_tokens=512)
+                code_response = luna_generate(enhanced_prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
                 return f"Luna (code + web-enhanced):\n{code_response.strip()}\n\n(Source: {url_used})"
-
         return f"Luna (code):\n{code_response.strip()}"
 
-    # Step 1: Try base Luna
+    # Base Luna response
     base_prompt = f"User: {user_question}\nLuna:"
-    base_answer = luna_generate(base_prompt)
+    base_answer = luna_generate(base_prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
 
     if any(kw in base_answer.lower() for kw in ["i don't know", "not sure", "as of", "unknown"]) or len(base_answer) < 20:
-        # Step 2: Try Wikipedia
         wiki_summary, wiki_url = try_wikipedia(user_question)
         if wiki_summary:
             prompt = f'''Use the following Wikipedia info to answer:
@@ -113,13 +109,11 @@ Context: {wiki_summary}
 
 Question: {user_question}
 Answer:'''
-            wiki_answer = luna_generate(prompt)
+            wiki_answer = luna_generate(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
             return f"Luna (Wikipedia): {wiki_answer.strip()}\n\n(Source: {wiki_url or 'Wikipedia'})"
 
-        # Step 3: Try web scraping fallback
         urls = get_top_webpages(user_question)
         context, url_used = scrape_first_good_content(urls)
-
         if context:
             prompt = f'''Use the following context to answer:
 
@@ -127,19 +121,9 @@ Context: {context}
 
 Question: {user_question}
 Answer:'''
-            web_answer = luna_generate(prompt)
+            web_answer = luna_generate(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
             return f"Luna (web-enhanced): {web_answer.strip()}\n\n(Source: {url_used})"
 
         return "Luna: I couldn’t find good info online right now."
 
     return f"Luna: {base_answer.strip()}"
-
-# ✅ Chat loop
-print("Chat with Luna (type 'exit' to quit):")
-while True:
-    user_input = input("You: ")
-    if user_input.lower() in ["exit", "quit"]:
-        print("Luna: Bye! See you again.")
-        break
-    response = smart_luna_answer(user_input)
-    print(response)