Update app.py

app.py

@@ -1,3 +1,198 @@
+import os
+import torch
+import gradio as gr
+from PyPDF2 import PdfReader
+from transformers import (
+    AutoTokenizer, pipeline,
+    AutoModelForCausalLM, AutoConfig,
+    BitsAndBytesConfig
+)
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.schema import Document
+from langchain import HuggingFacePipeline
+
+# ------------------------------
+# Device setup
+# ------------------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ------------------------------
+# Embedding model config
+# ------------------------------
+modelPath = "sentence-transformers/all-mpnet-base-v2"
+model_kwargs = {"device": str(device)}
+encode_kwargs = {"normalize_embedding": False}
+
+embeddings = HuggingFaceEmbeddings(
+    model_name=modelPath,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs
+)
+
+# ------------------------------
+# Load Mistral model in 4bit
+# ------------------------------
+model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+# 4-bit quantization config
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.float16
+)
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=bnb_config,
+    device_map="auto"
+)
+
+# ------------------------------
+# Improved Text Generation Pipeline
+# ------------------------------
+text_generation = pipeline(
+    model=model,
+    tokenizer=tokenizer,
+    task="text-generation",
+    temperature=0.7,
+    top_p=0.9,
+    top_k=50,
+    repetition_penalty=1.1,
+    return_full_text=False,
+    max_new_tokens=2000,
+    do_sample=True,
+    eos_token_id=tokenizer.eos_token_id,
+)
+
+# Wrap in LangChain interface
+mistral_llm = HuggingFacePipeline(pipeline=text_generation)
+
+# ------------------------------
+# PDF Processing Functions
+# ------------------------------
+def pdf_text(pdf_docs):
+    text = ""
+    for doc in pdf_docs:
+        reader = PdfReader(doc)
+        for page in reader.pages:
+            page_text = page.extract_text()
+            if page_text:
+                text += page_text + "\n"
+    return text
+
+def get_chunks(text):
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+        length_function=len
+    )
+    chunks = splitter.split_text(text)
+    return [Document(page_content=chunk) for chunk in chunks]
+
+def get_vectorstore(documents):
+    db = FAISS.from_documents(documents, embedding=embeddings)
+    db.save_local("faiss_index")
+
+# ------------------------------
+# Conversational Prompt Template
+# ------------------------------
+def get_qa_prompt():
+    prompt_template = """<s>[INST] 
+    You are a helpful, knowledgeable AI assistant. Answer the user's question based on the provided context.
+    
+    Guidelines:
+    - Respond in a natural, conversational tone
+    - Be detailed but concise
+    - Use paragraphs and bullet points when appropriate
+    - If you don't know, say so
+    - Maintain a friendly and professional demeanor
+    
+    Conversation History:
+    {chat_history}
+    
+    Relevant Context:
+    {context}
+    
+    Current Question: {question} 
+    
+    Provide a helpful response: [/INST]"""
+    
+    return PromptTemplate(
+        template=prompt_template,
+        input_variables=["context", "question", "chat_history"]
+    )
+
+# ------------------------------
+# Chat Handling Functions
+# ------------------------------
+def handle_pdf_upload(pdf_files):
+    try:
+        if not pdf_files:
+            return "⚠️ Please upload at least one PDF file"
+        
+        text = pdf_text(pdf_files)
+        if not text.strip():
+            return "⚠️ Could not extract text from PDFs - please try different files"
+            
+        chunks = get_chunks(text)
+        get_vectorstore(chunks)
+        return f"✅ Processed {len(pdf_files)} PDF(s) with {len(chunks)} text chunks"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def format_chat_history(chat_history):
+    return "\n".join([f"User: {q}\nAssistant: {a}" for q, a in chat_history[-3:]])
+
+def user_query(msg, chat_history):
+    if not os.path.exists("faiss_index"):
+        chat_history.append((msg, "Please upload PDF documents first so I can help you."))
+        return "", chat_history
+    
+    try:
+        # Load vector store
+        db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
+        retriever = db.as_retriever(search_kwargs={"k": 3})
+        
+        # Get relevant context
+        docs = retriever.get_relevant_documents(msg)
+        context = "\n\n".join([d.page_content for d in docs])
+        
+        # Generate response
+        prompt = get_qa_prompt()
+        chain = LLMChain(llm=mistral_llm, prompt=prompt)
+        
+        response = chain.run({
+            "question": msg,
+            "context": context,
+            "chat_history": format_chat_history(chat_history)
+        })
+        
+        # Clean response
+        response = response.strip()
+        for end_token in ["</s>", "[INST]", "[/INST]"]:
+            if response.endswith(end_token):
+                response = response[:-len(end_token)].strip()
+        
+        chat_history.append((msg, response))
+        return "", chat_history
+        
+    except Exception as e:
+        error_msg = f"Sorry, I encountered an error: {str(e)}"
+        chat_history.append((msg, error_msg))
+        return "", chat_history
+
+# ------------------------------
+# Gradio Interface
+# ------------------------------
 with gr.Blocks(theme=gr.themes.Soft(), title="PDF Chat Assistant") as demo:
     with gr.Row():
         gr.Markdown("""
@@ -79,4 +274,13 @@ with gr.Blocks(theme=gr.themes.Soft(), title="PDF Chat Assistant") as demo:
         None,
         chatbot,
         queue=False
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7861,
+        share=True,
+        debug=True
     )