Update app.py

app.py

@@ -1,118 +1,106 @@
-import gradio as gr
 import os
+import gradio as gr
 import torch
-from langchain_community.document_loaders import PyPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from huggingface_hub import login
+from langchain_community.document_loaders import PyMuPDFLoader, TextLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.chains import RetrievalQA
+from langchain_community.llms import HuggingFacePipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+
+# HF Authentication
+login(token=os.environ.get('HF_TOKEN'))
 
 # Configuration
-DOCS_DIR = "business_docs"
-EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+DOCS_DIR = "study_materials"
 MODEL_NAME = "microsoft/phi-2"
+EMBEDDINGS_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+MAX_TOKENS = 300
+CHUNK_SIZE = 1000
 
-def initialize_system():
-    # Document processing
-    if not os.path.exists(DOCS_DIR):
-        raise FileNotFoundError(f"Missing {DOCS_DIR} folder")
-    
-    pdf_files = [os.path.join(DOCS_DIR, f) 
-                for f in os.listdir(DOCS_DIR) 
-                if f.endswith(".pdf")]
+def load_documents():
+    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)}")
+    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=1000,
-        chunk_overlap=200
+        chunk_size=CHUNK_SIZE,
+        chunk_overlap=200,
+        separators=["\n\n", "\n", " "]
     )
-    
-    texts = []
-    for pdf in pdf_files:
-        loader = PyPDFLoader(pdf)
-        pages = loader.load_and_split(text_splitter)
-        texts.extend(pages)
+    texts = text_splitter.split_documents(documents)
 
-    # Create embeddings
-    embeddings = HuggingFaceEmbeddings(
-        model_name=EMBEDDING_MODEL,
-        model_kwargs={'device': 'cpu'}
-    )
-    
-    # Vector store
-    vector_store = FAISS.from_documents(texts, embeddings)
+    # Create vector store
+    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL)
+    db = FAISS.from_documents(texts, embeddings)
 
-    # Load model and tokenizer
+    # Load Phi-2 with authentication
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-    tokenizer.pad_token = tokenizer.eos_token  # Fix padding issue
-    
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
+        use_auth_token=True,  # Critical change for gated models
+        torch_dtype=torch.float32,
         trust_remote_code=True,
-        torch_dtype=torch.float32 if not torch.cuda.is_available() else torch.float16,
         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 vector_store, model, tokenizer
+    return RetrievalQA.from_chain_type(
+        llm=HuggingFacePipeline(pipeline=pipe),
+        chain_type="stuff",
+        retriever=db.as_retriever(search_kwargs={"k": 2}),
+        return_source_documents=True
+    )
 
-try:
-    vector_store, model, tokenizer = initialize_system()
-    print("✅ System initialized successfully")
-    if torch.cuda.is_available():
-        print("🚀 Using CUDA")
-        print(f"Memory usage: {torch.cuda.memory_allocated()/1024**3:.2f} GB")
-    else:
-        print("🧠 Using CPU")
-except Exception as e:
-    print(f"❌ Initialization failed: {str(e)}")
-    raise
+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 generate_response(query):
+def process_query(question, history):
     try:
-        # Context retrieval
-        docs = vector_store.similarity_search(query, k=3)
-        context = "\n".join([d.page_content for d in docs])
-
-        # Prompt template optimized for Phi-2
-        prompt = f"""Context:
-{context}
-
-Question: {query}
-Answer:"""
-        
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(
-            inputs.input_ids,
-            max_new_tokens=300,
-            temperature=0.3,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-        
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response.split("Answer:")[-1].strip()
-    
+        qa = create_qa_system()
+        formatted_q = f"Instruct: {question}\nOutput:"
+        response = qa.invoke({"query": formatted_q})
+        return format_response(response)
     except Exception as e:
-        return "Sorry, an error occurred while generating a response."
-
-# Gradio UI
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🧠 Enterprise Customer Support Chatbot")
-    
-    chatbot = gr.Chatbot(height=500, label="Conversation")
-    
-    with gr.Row():
-        msg = gr.Textbox(placeholder="Ask about our services...", scale=7)
-        submit_btn = gr.Button("Send", variant="primary", scale=1)
+        print(f"Error: {str(e)}")
+        return f"⚠️ Error: {str(e)[:100]}"
 
+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])
 
-    def respond(message, history):
-        response = generate_response(message)
-        history.append((message, response))
-        return "", history
-
-    submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    msg.submit(process_query, [msg, chatbot], [msg, chatbot])
 
-demo.launch(server_port=7860)
+if __name__ == "__main__":
+    app.launch(server_name="0.0.0.0", server_port=7860)