Update app.py

app.py

@@ -1,106 +1,92 @@
 import os
 import gradio as gr
-import torch
-from huggingface_hub import login
 from langchain_community.document_loaders import PyMuPDFLoader, TextLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_text_splitters import CharacterTextSplitter
 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, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer
 
-# HF Authentication
-login(token=os.environ.get('HF_TOKEN'))
-
-# Configuration
-DOCS_DIR = "study_materials"
-MODEL_NAME = "microsoft/phi-2"
-EMBEDDINGS_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-MAX_TOKENS = 300
-CHUNK_SIZE = 1000
-
-def load_documents():
+def load_documents(file_path="study_materials"):
     documents = []
-    for filename in os.listdir(DOCS_DIR):
-        path = os.path.join(DOCS_DIR, filename)
-        try:
-            if filename.endswith(".pdf"):
-                documents.extend(PyMuPDFLoader(path).load())
-            elif filename.endswith(".txt"):
-                documents.extend(TextLoader(path).load())
-        except Exception as e:
-            print(f"Error loading {filename}: {str(e)}")
+    for filename in os.listdir(file_path):
+        path = os.path.join(file_path, filename)
+        if filename.endswith(".pdf"):
+            loader = PyMuPDFLoader(path)
+            documents.extend(loader.load())
+        elif filename.endswith(".txt"):
+            loader = TextLoader(path)
+            documents.extend(loader.load())
     return documents
 
 def create_qa_system():
-    # Load and split documents
-    documents = load_documents()
-    if not documents:
-        raise gr.Error("No documents found in 'study_materials' folder")
-    
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=CHUNK_SIZE,
-        chunk_overlap=200,
-        separators=["\n\n", "\n", " "]
-    )
-    texts = text_splitter.split_documents(documents)
-
-    # Create vector store
-    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL)
-    db = FAISS.from_documents(texts, embeddings)
-
-    # Load Phi-2 with authentication
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        use_auth_token=True,  # Critical change for gated models
-        torch_dtype=torch.float32,
-        trust_remote_code=True,
-        device_map="auto",
-        low_cpu_mem_usage=True
-    )
-    
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=MAX_TOKENS,
-        temperature=0.7,
-        do_sample=True,
-        top_k=40,
-        device_map="auto"
-    )
-
-    return RetrievalQA.from_chain_type(
-        llm=HuggingFacePipeline(pipeline=pipe),
-        chain_type="stuff",
-        retriever=db.as_retriever(search_kwargs={"k": 2}),
-        return_source_documents=True
-    )
-
-def format_response(response):
-    answer = response["result"].split("</s>")[0].split("\nOutput:")[-1].strip()
-    sources = list({os.path.basename(doc.metadata["source"]) for doc in response["source_documents"]})
-    return f"{answer}\n\n📚 Sources: {', '.join(sources)}"
-
-def process_query(question, history):
     try:
-        qa = create_qa_system()
-        formatted_q = f"Instruct: {question}\nOutput:"
-        response = qa.invoke({"query": formatted_q})
-        return format_response(response)
+        # Load 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
+        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")  # ←
+        pipe = pipeline(
+            "text2text-generation",
+            model="google/flan-t5-large",
+            max_length=600,
+            temperature=0.7,
+            tokenizer=tokenizer,
+            do_sample=True,
+            top_k=50,
+            device=-1
+        )
+        
+        # Wrap pipeline in LangChain component
+        llm = HuggingFacePipeline(pipeline=pipe)
+        
+        # Create QA chain
+        return RetrievalQA.from_llm(
+            llm=llm,
+            retriever=db.as_retriever(search_kwargs={"k": 3}),
+            return_source_documents=True
+        )
     except Exception as e:
-        print(f"Error: {str(e)}")
-        return f"⚠️ Error: {str(e)[:100]}"
+        raise gr.Error(f"Error: {str(e)}")
 
-with gr.Blocks(title="Phi-2 Study Assistant", theme=gr.themes.Soft()) as app:
-    gr.Markdown("## 📚 Phi-2 Study Assistant\nUpload study materials to 'study_materials' and ask questions!")
-    chatbot = gr.Chatbot(height=400)
-    msg = gr.Textbox(label="Your Question")
-    clear = gr.ClearButton([msg, chatbot])
+# Initialize system
+try:
+    qa = create_qa_system()
+except Exception as e:
+    print(f"Startup failed: {str(e)}")
+    raise
 
-    msg.submit(process_query, [msg, chatbot], [msg, chatbot])
+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']})
+        return f"{answer}\n\n📚 Sources: {', '.join(sources)}"
+    except Exception as e:
+        return f"Error: {str(e)[:150]}"
 
-if __name__ == "__main__":
-    app.launch(server_name="0.0.0.0", server_port=7860)
+gr.ChatInterface(
+    ask_question,
+    title="Study Assistant",
+    description="Upload PDF/TXT files in 'study_materials' folder and ask questions!",
+    theme="soft"
+).launch()