luna-mobile / lunacode.py
cosmosai471's picture
Update lunacode.py
fdbc4f9 verified
import torch
import requests
from bs4 import BeautifulSoup
from duckduckgo_search import DDGS
from transformers import AutoModelForCausalLM, AutoTokenizer
import wikipedia
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")
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)
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)
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.
Task: {question}
Requirements:
- Use functions where appropriate.
- Add comments explaining each part of the code.
- Follow best coding practices.
- Do not add extra text outside the code.
Code:
'''
def try_wikipedia(query):
try:
wikipedia.set_lang("en")
page = wikipedia.page(query, auto_suggest=False)
summary = wikipedia.summary(query, auto_suggest=False)
return summary[:2000], page.url
except:
try:
summary = wikipedia.summary(query, auto_suggest=True)
return summary[:2000], None
except:
return None, None
def get_top_webpages(query, max_results=3):
urls = []
with DDGS() as ddgs:
results = ddgs.text(f"current {query}", max_results=max_results)
for result in results:
urls.append(result["href"])
return urls
def scrape_first_good_content(urls):
for url in urls:
try:
res = requests.get(url, timeout=10)
soup = BeautifulSoup(res.text, 'html.parser')
text = soup.get_text(separator=" ", strip=True)
if len(text) > 300:
return text[:2000], url
except:
continue
return None, None
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)
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)
if context:
enhanced_prompt = f'''You are a helpful AI programmer. Use the following web content as reference.
Context: {context}
Task: {user_question}
Generate well-structured code with inline comments and clear function design. Only return the code.
Code:
'''
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()}"
base_prompt = f"User: {user_question}\nLuna:"
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:
wiki_summary, wiki_url = try_wikipedia(user_question)
if wiki_summary:
prompt = f'''Use the following Wikipedia info to answer:
Context: {wiki_summary}
Question: {user_question}
Answer:'''
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'})"
urls = get_top_webpages(user_question)
context, url_used = scrape_first_good_content(urls)
if context:
prompt = f'''Use the following context to answer:
Context: {context}
Question: {user_question}
Answer:'''
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()}"