import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import spaces
from duckduckgo_search import DDGS
import time
import torch
from datetime import datetime
# Initialize model and tokenizer
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
# Modified model loading for CPU
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="cpu", # Changed to CPU
low_cpu_mem_usage=True,
torch_dtype=torch.float32 # Changed to float32 for CPU
)
def get_web_results(query, max_results=5): # Increased to 5 for better context
"""Get web search results using DuckDuckGo"""
try:
with DDGS() as ddgs:
results = list(ddgs.text(query, max_results=max_results))
return [{
"title": result.get("title", ""),
"snippet": result["body"],
"url": result["href"],
"date": result.get("published", "")
} for result in results]
except Exception as e:
return []
def format_prompt(query, context):
"""Format the prompt with web context"""
current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
context_lines = '\n'.join([f'- [{res["title"]}]: {res["snippet"]}' for res in context])
return f"""You are an intelligent search assistant. Answer the user's query using the provided web context.
Current Time: {current_time}
Query: {query}
Web Context:
{context_lines}
Provide a detailed answer in markdown format. Include relevant information from sources and cite them using [1], [2], etc.
Answer:"""
def format_sources(web_results):
"""Format sources with more details"""
if not web_results:
return "
No sources available
"
sources_html = "
"
for i, res in enumerate(web_results, 1):
title = res["title"] or "Source"
date = f"{res['date']}" if res['date'] else ""
sources_html += f"""