Update app.py

app.py

@@ -1,10 +1,10 @@
 import streamlit as st
 import tempfile
 import logging
-import time
 from typing import List
+import torch
 from langchain_community.document_loaders import PyPDFLoader
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain_community.llms import HuggingFacePipeline
 from langchain.chains.summarize import load_summarize_chain
@@ -19,12 +19,13 @@ logger = logging.getLogger(__name__)
 
 # Constants
 EMBEDDING_MODEL = 'sentence-transformers/all-MiniLM-L6-v2'
-DEFAULT_MODEL = "llava-v1.6-mistral-7b-hf"
+DEFAULT_MODEL = "microsoft/phi-2"
 
-# Cache expiration time for models (adjust as needed)
-MODEL_CACHE_EXPIRATION = 3600
+# Check for GPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
-@st.cache_resource(ttl=MODEL_CACHE_EXPIRATION)
+@st.cache_resource
 def load_embeddings():
     """Load and cache the embedding model."""
     try:
@@ -34,15 +35,15 @@ def load_embeddings():
         st.error("Failed to load the embedding model. Please try again later.")
         return None
 
-@st.cache_resource(ttl=MODEL_CACHE_EXPIRATION)
+@st.cache_resource
 def load_llm(model_name):
     """Load and cache the language model."""
     try:
-        pipe = pipeline("text-generation", model=model_name, max_length=1024)
+        pipe = pipeline("text-generation", model=model_name, device=device, max_length=1024)
         return HuggingFacePipeline(pipeline=pipe)
     except Exception as e:
         logger.error(f"Failed to load LLM: {e}")
-        st.error(f"Failed to load the model {model_name}. Please try again.")
+        st.error(f"Failed to load the model {model_name}. Please try another model or check your internet connection.")
         return None
 
 def process_pdf(file) -> List[Document]: