Update app.py

app.py

@@ -11,6 +11,7 @@ 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
 
 # Access the Hugging Face token from the environment variable
 HF_TOKEN = os.getenv("HF_Token")
@@ -29,10 +30,11 @@ papers_path = "./papers"
 os.makedirs(download_path, exist_ok=True)
 os.makedirs(papers_path, exist_ok=True)
 
-# Load LLaMA 2
+# 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, device_map="auto", torch_dtype="auto")
+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):
@@ -50,10 +52,6 @@ def add_local_files(module):
     elif module == "paper":
         return [{"title": os.path.basename(file_path), "url": None, "file_path": file_path, "type": "paper"}]
 
-import os
-import re
-from yt_dlp import YoutubeDL
-
 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"
@@ -140,14 +138,13 @@ def fetch_from_arxiv(query="machine learning", max_results=2, output_dir="./pape
             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").to("cuda")
+    inputs = tokenizer(input_text, return_tensors="pt")
     outputs = model.generate(inputs["input_ids"], max_length=500, temperature=0.7)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response