Update app.py

app.py

@@ -1,32 +1,37 @@
-import os
 import streamlit as st
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.llms import HuggingFaceHub
+from langchain.llms import HuggingFacePipeline
 from langchain.schema import Document
-import requests
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+from sentence_transformers import SentenceTransformer
+import torch
 from io import BytesIO
 import fitz  # PyMuPDF
-from dotenv import load_dotenv
 
 # Set device based on GPU availability
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load environment variables from .env file
-load_dotenv()
+# Load embeddings with a smaller model and run on CPU
+embedding_model = "all-MiniLM-L6-v2"
+embeddings = HuggingFaceEmbeddings(model_name=embedding_model, model_kwargs={'device': 'cpu'})
 
-# Hugging Face API token should now be loaded from the .env file
-# Explicitly set the Hugging Face API token from the environment variable
-os.environ["HUGGINGFACEHUB_API_TOKEN"] = os.getenv("HUGGINGFACE_API_TOKEN")
+# Set up text generation model with PyTorch-compatible pipeline
+model_name = "google/flan-t5-small"  # Or use a smaller model if needed
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
 
-# Load embeddings with Hugging Face API
-embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
-embeddings = HuggingFaceEmbeddings(model_name=embedding_model)  # Removed api_key parameter
+# Create a text generation pipeline
+generator = pipeline(
+    "text2text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if device == "cuda" else -1,
+    model_kwargs={"max_length": 256, "temperature": 0.7}
+)
 
-# Set up the text generation model using Hugging Face Hub
-model_name = "google/flan-t5-small"  # Use a smaller model to reduce response time and cost
-llm = HuggingFaceHub(repo_id=model_name, huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"), model_kwargs={"max_length": 256, "temperature": 0.7})
+llm = HuggingFacePipeline(pipeline=generator)
 
 # Streamlit interface
 def main():
@@ -68,12 +73,12 @@ def main():
             # Concatenate retrieved docs into a single prompt
             prompt = "\n".join([doc.page_content for doc in docs]) + "\n\n" + user_input
 
-            # Generate response using the Hugging Face API
+            # Generate response
             try:
-                response = llm(prompt)
+                response = generator(prompt, max_new_tokens=50, num_return_sequences=1)[0]["generated_text"]
                 st.write(response)
-            except requests.exceptions.RequestException as e:
-                st.error(f"Error connecting to Hugging Face API: {e}")
+            except torch.cuda.OutOfMemoryError:
+                st.error("Out of memory. Try using a smaller model or fewer documents.")
 
 if __name__ == "__main__":
-    main()
+    main()