Update app.py

app.py

@@ -1,4 +1,6 @@
 import os
+import gc
+import torch
 import gradio as gr
 from langchain_community.document_loaders import PyMuPDFLoader, TextLoader
 from langchain_text_splitters import CharacterTextSplitter
@@ -6,14 +8,12 @@ from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
 from langchain_community.llms import HuggingFacePipeline
-from transformers import pipeline, AutoTokenizer
-
-# Optional but recommended addition
+from transformers import pipeline, AutoTokenizer, BitsAndBytesConfig
 from huggingface_hub import login
-import os
 
+# Handle HF token securely
 if os.environ.get("HF_TOKEN"):
-    login(token=os.environ["HF_TOKEN"])  # No hardcoded tokens!
+    login(token=os.environ["HF_TOKEN"])
 
 def load_documents(file_path="study_materials"):
     documents = []
@@ -29,46 +29,55 @@ def load_documents(file_path="study_materials"):
 
 def create_qa_system():
     try:
-        # Load documents
+        # Load and split documents
         documents = load_documents()
         if not documents:
             raise ValueError("📚 No study materials found")
-        
-        # Text splitting
+            
         text_splitter = CharacterTextSplitter(
             chunk_size=1100,
             chunk_overlap=200,
             separator="\n\n"
         )
         texts = text_splitter.split_documents(documents)
-        
-        # Embeddings
+
+        # Create embeddings
         embeddings = HuggingFaceEmbeddings(
             model_name="sentence-transformers/all-MiniLM-L6-v2"
         )
         
         # Vector store
         db = FAISS.from_documents(texts, embeddings)
-        
-        # LLM setup with proper LangChain wrapper
-        tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")  # ←
+
+        # Quantization config
+        quant_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+            llm_int8_threshold=6.0
+        )
+
+        # LLM setup with optimizations
+        tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
         pipe = pipeline(
             "text2text-generation",
             model="google/flan-t5-large",
-            max_length=600,
-            temperature=0.7,
             tokenizer=tokenizer,
+            max_length=400,
+            temperature=0.7,
             do_sample=True,
             top_k=50,
-            device=-1
+            device=-1,  # Force CPU usage
+            model_kwargs={
+                "torch_dtype": torch.float16,
+                "quantization_config": quant_config
+            }
         )
-        
-        # Wrap pipeline in LangChain component
-        llm = HuggingFacePipeline(pipeline=pipe)
-        
-        # Create QA chain
+
+        # Memory cleanup
+        gc.collect()
+
+        # Create QA system
         return RetrievalQA.from_llm(
-            llm=llm,
+            llm=HuggingFacePipeline(pipeline=pipe),
             retriever=db.as_retriever(search_kwargs={"k": 3}),
             return_source_documents=True
         )
@@ -86,7 +95,7 @@ def ask_question(question, history):
     try:
         result = qa.invoke({"query": question})
         answer = result["result"]
-        sources = list({doc.metadata['source'] for doc in result['source_documents']})
+        sources = {os.path.basename(doc.metadata['source']) for doc in result['source_documents']}
         return f"{answer}\n\n📚 Sources: {', '.join(sources)}"
     except Exception as e:
         return f"Error: {str(e)[:150]}"
@@ -96,4 +105,4 @@ gr.ChatInterface(
     title="Study Assistant",
     description="Upload PDF/TXT files in 'study_materials' folder and ask questions!",
     theme="soft"
-).launch()
+).launch()