app.py

# Import libraries
import os
import requests
import re
from yt_dlp import YoutubeDL
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import login
import arxiv
import numpy as np
import torch  # Add torch to explicitly set the device
import gradio as gr

# Access the Hugging Face token from the environment variable
HF_TOKEN = os.getenv("HF_Token")
login(token=HF_TOKEN)

# Initialize the embedding model
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

# Define paths for downloaded content and database
file_paths = {
    "video": "./Machine Learning.mp4",  # Replace with actual paths
    "paper": "./DeepSeek_v3.pdf",
}
download_path = "./downloads"
papers_path = "./papers"
os.makedirs(download_path, exist_ok=True)
os.makedirs(papers_path, exist_ok=True)

# Load LLaMA 2 (set to use CPU)
model_name = "meta-llama/Llama-3.2-1B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)  # Ensure float32 for CPU
model.to("cpu")  # Explicitly set the model to use the CPU

# Define utility functions
def compute_similarity(query_embedding, content_embeddings):
    """Compute cosine similarity between query and content embeddings."""
    similarities = cosine_similarity([query_embedding], content_embeddings).flatten()
    return similarities

def add_local_files(module):
    """Add local files from the database to the metadata."""
    if module not in file_paths:
        return []
    file_path = file_paths[module]
    if module == "video":
        return [{"title": os.path.basename(file_path), "url": None, "file_path": file_path, "type": "video"}]
    elif module == "paper":
        return [{"title": os.path.basename(file_path), "url": None, "file_path": file_path, "type": "paper"}]

def download_youtube_video(video_url, output_dir, title=None):
    """Download a YouTube video using yt_dlp."""
    sanitized_title = re.sub(r'[\\/*?:"<>|]', '_', title) if title else "unknown_title"
    ydl_opts = {
        'quiet': True,
        'outtmpl': f"{output_dir}/{sanitized_title}.%(ext)s",
        'format': 'best',
    }
    try:
        with YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(video_url, download=True)
            downloaded_file = ydl.prepare_filename(info)
        return downloaded_file
    except Exception as e:
        print(f"Failed to download video {video_url}. Error: {e}")
        return None

def fetch_and_download_youtube_video(query, output_dir="./videos"):
    """Fetch and download a YouTube video based on a query."""
    print(f"Fetching YouTube video for query: '{query}'")
    ydl_opts = {
        'quiet': True,
        'format': 'best',
        'outtmpl': f"{output_dir}/%(title)s.%(ext)s",  # Default template
    }

    try:
        with YoutubeDL(ydl_opts) as ydl:
            search_results = ydl.extract_info(f"ytsearch:{query}", download=False)
            if 'entries' not in search_results or len(search_results['entries']) == 0:
                print(f"No YouTube results found for query: '{query}'")
                return []

            video_info = search_results['entries'][0]
            video_title = video_info.get("title", "unknown_title")
            video_url = video_info.get("webpage_url", None)
            if not video_url:
                print("No URL found for the video.")
                return []

            local_path = download_youtube_video(video_url, output_dir, title=video_title)
            if not local_path:
                return []

            print(f"Successfully downloaded video: {video_title}")
            return [{"title": video_title, "url": video_url, "file_path": local_path, "type": "video"}]

    except Exception as e:
        print(f"Error fetching YouTube video for query '{query}': {e}")
        return []

def fetch_from_arxiv(query="machine learning", max_results=2, output_dir="./papers"):
    """Fetch papers from arXiv and download their PDFs."""
    print(f"Fetching papers for query: {query}")
    client = arxiv.Client()
    search = arxiv.Search(
        query=query,
        max_results=max_results,
        sort_by=arxiv.SortCriterion.Relevance
    )
    metadata = []
    for i, result in enumerate(client.results(search)):
        pdf_url = result.pdf_url
        filename = f"{query.replace(' ', '_')}_arxiv_{i}.pdf"
        local_path = os.path.join(output_dir, filename)
        try:
            response = requests.get(pdf_url)
            if response.status_code == 200:
                with open(local_path, 'wb') as f:
                    f.write(response.content)
                print(f"Downloaded paper: {filename}")
                metadata.append({"title": result.title, "url": pdf_url, "file_path": local_path, "type": "paper"})
            else:
                print(f"Failed to download paper: {pdf_url}. Status code: {response.status_code}")
        except Exception as e:
            print(f"Error downloading paper: {e}")
    return metadata

def generate_llama_response(query, context=None):
    """Generate a response using LLaMA 2."""
    input_text = f"Query: {query}\n"
    if context:
        input_text += f"Context: {context}\n"
    input_text += "Answer:"
    inputs = tokenizer(input_text, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_length=40, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

def hybrid_rag_system_with_llama(query):
    """Use LLaMA 2 to generate a final response after retrieving the best video and paper."""
    modules = ["video", "paper"]
    final_results = {}
    query_embedding = embedding_model.encode(query)

    for module in modules:
        metadata = []
        metadata.extend(add_local_files(module))
        if module == "video":
            metadata.extend(fetch_and_download_youtube_video(query, output_dir=download_path))
        elif module == "paper":
            metadata.extend(fetch_from_arxiv(query, max_results=2, output_dir=papers_path))
        if metadata:
            descriptions = [f"{item['title']} ({item['type']})" for item in metadata]
            description_embeddings = [embedding_model.encode(description) for description in descriptions]
            similarities = compute_similarity(query_embedding, description_embeddings)
            for idx, item in enumerate(metadata):
                item["similarity"] = similarities[idx]
            best_match_idx = np.argmax(similarities)
            final_results[module] = {
                "best_match": metadata[best_match_idx],
                "similarity": similarities[best_match_idx],
                "all_metadata": metadata,
            }
        else:
            final_results[module] = {"best_match": None, "similarity": None, "all_metadata": []}
    video_context = f"Best Video: {final_results['video']['best_match']['title']}" if final_results['video']['best_match'] else "No relevant video found."
    paper_context = f"Best Paper: {final_results['paper']['best_match']['title']}" if final_results['paper']['best_match'] else "No relevant paper found."
    context = f"{video_context}\n{paper_context}"
    final_response = generate_llama_response(query, context)
    return final_results, final_response

# Define Gradio interface
def gradio_interface(query):
    """Gradio wrapper for hybrid RAG system."""
    _, final_response = hybrid_rag_system_with_llama(query)
    return final_response

# Create Gradio app
interface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(label="Enter your query", placeholder="e.g., short easy machine learning"),
    outputs=gr.Textbox(label="Generated Response"),
    title="Hybrid RAG System with LLaMA",
    description="Enter a query to retrieve relevant resources and generate a response using LLaMA."
)

# Launch Gradio app
if __name__ == "__main__":
    interface.launch()