Create app.py

app.py

@@ -0,0 +1,138 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from sentence_transformers import SentenceTransformer
+from PyPDF2 import PdfReader
+import numpy as np
+import torch
+
+class RAGChatbot:
+    def __init__(self, model_name="TheBloke/Mistral-7B-Instruct-v0.1-GPTQ", embedding_model="all-MiniLM-L6-v2"):
+        # Initialize tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+        
+        # Initialize embedding model
+        self.embedding_model = SentenceTransformer(embedding_model)
+        
+        # Initialize document storage
+        self.documents = []
+        self.embeddings = []
+
+    def extract_text_from_pdf(self, pdf_path):
+        reader = PdfReader(pdf_path)
+        text = ""
+        for page in reader.pages:
+            text += page.extract_text() + "\n"
+        return text
+
+    def load_documents(self, file_paths):
+        self.documents = []
+        self.embeddings = []
+        
+        for file_path in file_paths:
+            if file_path.endswith('.pdf'):
+                text = self.extract_text_from_pdf(file_path)
+            else:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    text = f.read()
+            
+            # Split text into chunks
+            chunks = [text[i:i+500] for i in range(0, len(text), 500)]
+            self.documents.extend(chunks)
+        
+        # Generate embeddings
+        self.embeddings = self.embedding_model.encode(self.documents)
+        return f"Loaded {len(self.documents)} text chunks from {len(file_paths)} files"
+
+    def retrieve_relevant_context(self, query, top_k=3):
+        if not self.documents:
+            return "No documents loaded"
+        
+        # Generate query embedding
+        query_embedding = self.embedding_model.encode([query])[0]
+        
+        # Calculate cosine similarities
+        similarities = np.dot(self.embeddings, query_embedding) / (
+            np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_embedding)
+        )
+        
+        # Get top k most similar documents
+        top_indices = similarities.argsort()[-top_k:][::-1]
+        return " ".join([self.documents[i] for i in top_indices])
+
+    def generate_response(self, query, context):
+        # Construct prompt with context
+        full_prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
+        
+        # Generate response
+        inputs = self.tokenizer(full_prompt, return_tensors="pt").to(self.model.device)
+        outputs = self.model.generate(**inputs, max_new_tokens=150)
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return response.split("Answer:")[-1].strip()
+
+    def chat(self, query, history):
+        try:
+            # Retrieve relevant context
+            context = self.retrieve_relevant_context(query)
+            
+            # Generate response
+            response = self.generate_response(query, context)
+            
+            return response
+        except Exception as e:
+            return f"An error occurred: {str(e)}"
+
+# Create Gradio interface
+def create_interface():
+    rag_chatbot = RAGChatbot()
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# RAG Chatbot with Hugging Face Models")
+        
+        with gr.Row():
+            file_input = gr.File(label="Upload Documents", file_count="multiple", type="filepath")
+            load_btn = gr.Button("Load Documents")
+        
+        status_output = gr.Textbox(label="Load Status")
+        
+        chatbot = gr.Chatbot()
+        msg = gr.Textbox(label="Enter your query")
+        submit_btn = gr.Button("Send")
+        clear_btn = gr.Button("Clear Chat")
+
+        # Event handlers
+        load_btn.click(
+            rag_chatbot.load_documents, 
+            inputs=[file_input], 
+            outputs=[status_output]
+        )
+        
+        submit_btn.click(
+            rag_chatbot.chat, 
+            inputs=[msg, chatbot], 
+            outputs=[chatbot]
+        ).then(
+            lambda: gr.Textbox(interactive=True), 
+            None, 
+            [msg]
+        )
+        
+        msg.submit(
+            rag_chatbot.chat, 
+            inputs=[msg, chatbot], 
+            outputs=[chatbot]
+        ).then(
+            lambda: gr.Textbox(interactive=True), 
+            None, 
+            [msg]
+        )
+        
+        clear_btn.click(lambda: None, None, [chatbot, msg])
+
+    return demo
+
+# Launch the app
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()