import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import spaces
from duckduckgo_search import DDGS
import time
import torch
from datetime import datetime
import os
import subprocess
import numpy as np
# Install required dependencies for Kokoro with better error handling
try:
subprocess.run(['git', 'lfs', 'install'], check=True)
if not os.path.exists('Kokoro-82M'):
subprocess.run(['git', 'clone', 'https://huggingface.co/hexgrad/Kokoro-82M'], check=True)
# Try installing espeak with proper package manager commands
try:
# Update package list first
subprocess.run(['apt-get', 'update'], check=True)
# Try installing espeak first (more widely available)
subprocess.run(['apt-get', 'install', '-y', 'espeak'], check=True)
except subprocess.CalledProcessError:
print("Warning: Could not install espeak. Attempting espeak-ng...")
try:
subprocess.run(['apt-get', 'install', '-y', 'espeak-ng'], check=True)
except subprocess.CalledProcessError:
print("Warning: Could not install espeak or espeak-ng. TTS functionality may be limited.")
except Exception as e:
print(f"Warning: Initial setup error: {str(e)}")
print("Continuing with limited functionality...")
# Initialize models and tokenizers
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
# Move model initialization inside a function to prevent CUDA initialization in main process
def init_models():
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
offload_folder="offload",
low_cpu_mem_usage=True,
torch_dtype=torch.float16
)
return model
# Initialize Kokoro TTS with better error handling
try:
import sys
sys.path.append('Kokoro-82M')
from models import build_model
from kokoro import generate
# Don't initialize models/voices in main process for ZeroGPU compatibility
VOICE_CHOICES = {
'πΊπΈ Female (Default)': 'af',
'πΊπΈ Bella': 'af_bella',
'πΊπΈ Sarah': 'af_sarah',
'πΊπΈ Nicole': 'af_nicole'
}
TTS_ENABLED = True
except Exception as e:
print(f"Warning: Could not initialize Kokoro TTS: {str(e)}")
TTS_ENABLED = False
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}
Important: For election-related queries, please distinguish clearly between different election years and types (presidential vs. non-presidential). Only use information from the provided web context.
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. If the query is about elections, clearly specify which year and type of election you're discussing.
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"""