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
from typing import List, Dict, Tuple, Any
# 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...")
# --- Initialization (Do this ONCE) ---
model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
# Initialize DeepSeek model
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
offload_folder="offload",
low_cpu_mem_usage=True,
torch_dtype=torch.float16
)
# Initialize Kokoro TTS (with error handling)
VOICE_CHOICES = {
'πΊπΈ Female (Default)': 'af',
'πΊπΈ Bella': 'af_bella',
'πΊπΈ Sarah': 'af_sarah',
'πΊπΈ Nicole': 'af_nicole'
}
TTS_ENABLED = False
TTS_MODEL = None
VOICEPACK = None
try:
if os.path.exists('Kokoro-82M'):
import sys
sys.path.append('Kokoro-82M')
from models import build_model # type: ignore
from kokoro import generate # type: ignore
device = 'cuda' if torch.cuda.is_available() else 'cpu' # Correct device handling
TTS_MODEL = build_model('Kokoro-82M/kokoro-v0_19.pth', device)
# Load default voice
try:
VOICEPACK = torch.load('Kokoro-82M/voices/af.pt', map_location=device, weights_only=True)
except Exception as e:
print(f"Warning: Could not load default voice: {e}")
raise
TTS_ENABLED = True
else:
print("Warning: Kokoro-82M directory not found. TTS disabled.")
except Exception as e:
print(f"Warning: Could not initialize Kokoro TTS: {str(e)}")
TTS_ENABLED = False
def get_web_results(query: str, max_results: int = 5) -> List[Dict[str, str]]:
"""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:
print(f"Error in web search: {e}")
return []
def format_prompt(query: str, context: List[Dict[str, str]]) -> str:
"""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: List[Dict[str, str]]) -> str:
"""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"""