Update app.py

app.py

@@ -1,16 +1,44 @@
 import gradio as gr
+import fitz  # PyMuPDF for PDF text extraction
+import faiss  # FAISS for vector search
+import numpy as np
+from sentence_transformers import SentenceTransformer
 from huggingface_hub import InferenceClient
 from typing import List, Tuple
 
 # Default settings
 class ChatConfig:
     MODEL = "google/gemma-3-27b-it"
-    DEFAULT_SYSTEM_MSG = "You are a super intelligent and useful Chatbot."
+    DEFAULT_SYSTEM_MSG = "You are an AI assistant answering only based on the uploaded PDF."
     DEFAULT_MAX_TOKENS = 512
     DEFAULT_TEMP = 0.3
     DEFAULT_TOP_P = 0.95
 
 client = InferenceClient(ChatConfig.MODEL)
+embed_model = SentenceTransformer("all-MiniLM-L6-v2")  # Lightweight embedding model
+vector_dim = 384  # Embedding size
+index = faiss.IndexFlatL2(vector_dim)  # FAISS index
+
+documents = []  # Store extracted text
+
+def extract_text_from_pdf(pdf_path):
+    """Extracts text from PDF"""
+    doc = fitz.open(pdf_path)
+    text_chunks = [page.get_text("text") for page in doc]
+    return text_chunks
+
+def create_vector_db(text_chunks):
+    """Embeds text chunks and adds them to FAISS index"""
+    global documents, index
+    documents = text_chunks
+    embeddings = embed_model.encode(text_chunks)
+    index.add(np.array(embeddings, dtype=np.float32))
+
+def search_relevant_text(query):
+    """Finds the most relevant text chunk for the given query"""
+    query_embedding = embed_model.encode([query])
+    _, closest_idx = index.search(np.array(query_embedding, dtype=np.float32), k=3)
+    return "\n".join([documents[i] for i in closest_idx[0]])
 
 def generate_response(
     message: str,
@@ -20,17 +48,20 @@ def generate_response(
     temperature: float = ChatConfig.DEFAULT_TEMP,
     top_p: float = ChatConfig.DEFAULT_TOP_P
 ) -> str:
+    if not documents:
+        return "Please upload a PDF first."
+
+    context = search_relevant_text(message)  # Get relevant content from PDF
+
     messages = [{"role": "system", "content": system_message}]
-    
-    # Conversation history
     for user_msg, bot_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
         if bot_msg:
             messages.append({"role": "assistant", "content": bot_msg})
-    
-    messages.append({"role": "user", "content": message})
-    
+
+    messages.append({"role": "user", "content": f"Context: {context}\nQuestion: {message}"})
+
     response = ""
     for chunk in client.chat_completion(
         messages,
@@ -43,90 +74,39 @@ def generate_response(
         response += token
         yield response
 
+def handle_upload(pdf_file):
+    """Handles PDF upload and creates vector DB"""
+    text_chunks = extract_text_from_pdf(pdf_file.name)
+    create_vector_db(text_chunks)
+    return "PDF uploaded and indexed successfully!"
 
-def create_interface() -> gr.ChatInterface:
-    """Create and configure the chat interface."""
-    # Custom CSS for a modern look
-    custom_css = """
-    .chatbot .message { 
-        border-radius: 12px; 
-        margin: 5px; 
-        padding: 10px; 
-    }
-    .chatbot .user-message { 
-        background-color: #e3f2fd; 
-    }
-    .chatbot .bot-message { 
-        background-color: #f5f5f5; 
-    }
-    .gr-button { 
-        border-radius: 8px; 
-        padding: 8px 16px; 
-    }
-    """
-
-    # Custom chatbot 
-    chatbot = gr.Chatbot(
-        label="Gemma Chat",
-        avatar_images=("./user.png", "./botge.png"),
-        height=450,
-        show_copy_button=True
-    )
-
-    # Chat interface
-    interface = gr.ChatInterface(
-        fn=generate_response,
-        chatbot=chatbot,
-        title="欢迎体验 喵哥 Google-Gemma-3大模型",
-        theme=gr.themes.Soft(),
-        css=custom_css,
-        additional_inputs=[
-            gr.Textbox(
-                value=ChatConfig.DEFAULT_SYSTEM_MSG,
-                label="系统提示词",
-                lines=2,
-                placeholder="Enter system message..."
-            ),
-            gr.Slider(
-                minimum=1,
-                maximum=8192,
-                value=ChatConfig.DEFAULT_MAX_TOKENS,
-                step=1,
-                label="Max Tokens",
-                info="Controls response length"
-            ),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=ChatConfig.DEFAULT_TEMP,
-                step=0.1,
-                label="Temperature",
-                info="Controls randomness"
-            ),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=ChatConfig.DEFAULT_TOP_P,
-                step=0.05,
-                label="Top-P",
-                info="Controls diversity"
-            )
-        ],
-        additional_inputs_accordion=gr.Accordion(label="高级设置", open=False)
-    )
-    
-    return interface
+def create_interface() -> gr.Blocks:
+    """Creates the Gradio interface"""
+    with gr.Blocks() as interface:
+        gr.Markdown("# PDF-Based Chatbot using Google Gemma")
+        
+        with gr.Row():
+            chatbot = gr.Chatbot(label="Chat with Your PDF")
+            pdf_upload = gr.File(label="Upload PDF", type="file")
+        
+        with gr.Row():
+            user_input = gr.Textbox(label="Ask a question", placeholder="Type here...")
+            send_button = gr.Button("Send")
+        
+        output = gr.Textbox(label="Response", lines=5)
 
-def main():
-    app = create_interface()
-    app.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_api=False,
-        show_error=True,
-        debug=True
-    )
+        # Upload PDF handler
+        pdf_upload.change(handle_upload, inputs=[pdf_upload], outputs=[])
+
+        # Chat function
+        send_button.click(
+            generate_response, 
+            inputs=[user_input, chatbot], 
+            outputs=[output]
+        )
+
+    return interface
 
 if __name__ == "__main__":
-    main()
+    app = create_interface()
+    app.launch()