Spaces:
Running
Running
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()}" | |