import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import spaces
from duckduckgo_search import DDGS
import time
import torch
from datetime import datetime
import gc # For manual garbage collection
# Initialize model and tokenizer with optimizations
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
# Load config first to set optimal parameters
config = AutoConfig.from_pretrained(model_name)
config.use_cache = True # Enable KV-caching for faster inference
# Initialize tokenizer with optimizations
tokenizer = AutoTokenizer.from_pretrained(
model_name,
model_max_length=256, # Reduced for faster processing
padding_side="left",
truncation_side="left",
)
tokenizer.pad_token = tokenizer.eos_token
# Load model with optimizations
model = AutoModelForCausalLM.from_pretrained(
model_name,
config=config,
device_map="cpu",
low_cpu_mem_usage=True,
torch_dtype=torch.float32,
)
# Enable model optimizations
model.eval() # Set to evaluation mode
torch.set_num_threads(4) # Limit CPU threads for better performance
def get_web_results(query, max_results=3): # Reduced max results
"""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"][:200], # Limit snippet length
"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 - optimized version"""
context_lines = '\n'.join([f'[{i+1}] {res["snippet"]}'
for i, res in enumerate(context)])
return f"""Answer this query using the context: {query}\n\nContext:\n{context_lines}\n\nAnswer:"""
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"""