Update app.py

app.py

@@ -2,8 +2,8 @@ import os
 import gradio as gr
 import numpy as np
 import google.generativeai as genai
-from sentence_transformers import SentenceTransformer
 import faiss
+from sentence_transformers import SentenceTransformer
 from datasets import load_dataset
 from dotenv import load_dotenv
 
@@ -13,50 +13,44 @@ load_dotenv()
 # Configuration
 MODEL_NAME = "all-MiniLM-L6-v2"
 GENAI_MODEL = "gemini-pro"
-DATASET_LINK = "curl -X GET \
-     "https://datasets-server.huggingface.co/rows?dataset=midrees2806%2F7K_Dataset&config=default&split=train&offset=100&length=100""
+DATASET_NAME = "midrees2806/7K_Dataset"  # Direct dataset name
 CHUNK_SIZE = 500
 TOP_K = 3
 
-# Initialize models
-try:
-    embedding_model = SentenceTransformer(MODEL_NAME)
-except Exception as e:
-    raise RuntimeError(f"Failed to initialize embedding model: {str(e)}")
-
 class GeminiRAGSystem:
     def __init__(self):
         self.index = None
         self.chunks = []
         self.dataset_loaded = False
         self.gemini_api_key = os.getenv("AIzaSyASrFvE3gFPigihza0JTuALzZmBx0Kc3d0")
+        
+        # Initialize embedding model
+        try:
+            self.embedding_model = SentenceTransformer(MODEL_NAME)
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize embedding model: {str(e)}")
+        
+        # Configure Gemini
         if self.gemini_api_key:
             genai.configure(api_key=self.gemini_api_key)
     
-    def load_dataset_from_link(self, dataset_link: str):
-        """Load dataset from Hugging Face link"""
+    def load_dataset(self):
+        """Load dataset from Hugging Face"""
         try:
-            # Extract dataset name from URL
-            dataset_name = dataset_link.split("datasets/")[-1].split("/")[0].strip()
-            if not dataset_name:
-                raise ValueError("Invalid dataset URL format")
-            
             with gr.Progress() as progress:
                 progress(0.1, desc="📦 Downloading dataset...")
-                dataset = load_dataset(dataset_name, split='train')
+                dataset = load_dataset(DATASET_NAME, split='train')
                 
                 progress(0.5, desc="🔨 Processing dataset...")
-                if 'context' in dataset.features:
-                    self.chunks = list(set(dataset['context']))
-                elif 'text' in dataset.features:
-                    self.chunks = dataset['text']
-                elif 'question' in dataset.features and 'answer' in dataset.features:
-                    self.chunks = [f"Q: {q}\nA: {a}" for q, a in zip(dataset['question'], dataset['answer'])]
+                if 'text' in dataset.features:
+                    self.chunks = dataset['text'][:1000]  # Limit to first 1000 entries
+                elif 'context' in dataset.features:
+                    self.chunks = dataset['context'][:1000]
                 else:
-                    raise ValueError("Unsupported dataset format")
+                    raise ValueError("Dataset must have 'text' or 'context' field")
                 
                 progress(0.7, desc="🧠 Creating embeddings...")
-                embeddings = embedding_model.encode(self.chunks, show_progress_bar=False)
+                embeddings = self.embedding_model.encode(self.chunks, show_progress_bar=False)
                 self.index = faiss.IndexFlatL2(embeddings.shape[1])
                 self.index.add(embeddings.astype('float32'))
                 
@@ -68,168 +62,74 @@ class GeminiRAGSystem:
             return False
     
     def get_relevant_context(self, query: str) -> str:
-        """Retrieve most relevant chunks with scores"""
-        if not self.index or not self.chunks:
+        """Retrieve most relevant chunks"""
+        if not self.index:
             return ""
             
-        query_embed = embedding_model.encode([query])
-        scores, indices = self.index.search(query_embed.astype('float32'), k=TOP_K)
+        query_embed = self.embedding_model.encode([query])
+        _, indices = self.index.search(query_embed.astype('float32'), k=TOP_K)
         
         context = []
-        for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
+        for idx in indices[0]:
             if idx < len(self.chunks):
-                context.append(f"🔍 Match {i+1} (Score: {1-score:.2f}):\n{self.chunks[idx]}\n")
-        return "\n".join(context)
+                context.append(self.chunks[idx])
+        return "\n\n".join(context)
     
     def generate_response(self, query: str) -> str:
-        """Generate response using only dataset context"""
+        """Generate response using Gemini"""
         if not self.dataset_loaded:
             return "⚠️ Please load the dataset first"
         if not self.gemini_api_key:
-            return "🔑 Please set your Gemini API key in environment variables"
+            return "🔑 Please set your Gemini API key"
         
         context = self.get_relevant_context(query)
         if not context:
-            return "No relevant context found in dataset"
+            return "No relevant context found"
         
-        prompt = f"""You are an expert AI assistant that answers STRICTLY based on the provided context.
-        Follow these rules:
-        1. Answer concisely using ONLY the context below
-        2. If the answer isn't in the context, say "I couldn't find this in the dataset"
-        3. Never make up information
-        4. For ambiguous questions, ask for clarification
-        Context:
+        prompt = f"""Answer based on this context:
         {context}
+        
         Question: {query}
-        Answer:"""
+        Answer concisely:"""
         
         try:
             model = genai.GenerativeModel(GENAI_MODEL)
             response = model.generate_content(prompt)
             return response.text
         except Exception as e:
-            return f"⚠️ Error generating response: {str(e)}"
+            return f"⚠️ Error: {str(e)}"
 
-# Initialize the RAG system
+# Initialize system
 rag_system = GeminiRAGSystem()
 
-# Custom CSS for modern UI
-css = """
-.gradio-container {
-    max-width: 900px !important;
-    margin: auto !important;
-    font-family: 'Inter', sans-serif;
-}
-.dark .gradio-container {
-    background-color: #1e1e2e;
-}
-.message-user {
-    background: #3b82f6;
-    color: white;
-    border-radius: 18px 18px 0 18px;
-    padding: 12px;
-    margin: 8px 0;
-    max-width: 80%;
-    margin-left: auto;
-}
-.message-bot {
-    background: #f3f4f6;
-    color: #111827;
-    border-radius: 18px 18px 18px 0;
-    padding: 12px;
-    margin: 8px 0;
-    max-width: 80%;
-}
-.dark .message-bot {
-    background: #2d3748;
-    color: #f7fafc;
-}
-.progress-bar {
-    height: 6px !important;
-}
-"""
-
-# Chat interface
-with gr.Blocks(css=css, theme=gr.themes.Default()) as app:
-    # Store chat history
-    chat_history = gr.State([])
-    
-    gr.Markdown("## UE-ChatBot")
-    gr.Markdown(f"**Dataset:** {DATASET_LINK}")
+# Create interface
+with gr.Blocks(title="RAG Chatbot") as app:
+    gr.Markdown("# UE_ChatBot")
     
     with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### ⚙️ Configuration")
-            dataset_url = gr.Textbox(
-                label="Hugging Face Dataset URL",
-                value=DATASET_LINK,
-                interactive=True
-            )
+        with gr.Column():
             load_btn = gr.Button("🚀 Load Dataset", variant="primary")
-            status = gr.Markdown("ℹ️ Please load the dataset first")
+            status = gr.Markdown("ℹ️ Click to load dataset")
         
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(
-                label="Chat History",
-                bubble_full_width=False,
-                avatar_images=(
-                    "https://avatars.githubusercontent.com/u/1561194?v=4",  # User avatar
-                    "https://huggingface.co/spaces/groq/Groq-LLM/resolve/main/groq_logo.png"  # Bot avatar
-                )
-            )
-            query = gr.Textbox(
-                label="Type your question...",
-                placeholder="Ask about the dataset content",
-                autofocus=True
-            )
-            with gr.Row():
-                submit_btn = gr.Button("📤 Submit", variant="primary")
-                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
+        with gr.Column():
+            chatbot = gr.Chatbot()
+            query = gr.Textbox(label="Your question", placeholder="Ask about the dataset...")
+            submit_btn = gr.Button("📤 Submit", variant="primary")
     
     # Event handlers
-    def load_dataset(dataset_url):
-        if rag_system.load_dataset_from_link(dataset_url):
-            return "✅ Dataset loaded successfully!"
+    def load_dataset():
+        if rag_system.load_dataset():
+            return "✅ Dataset ready! You can now ask questions."
         return "❌ Failed to load dataset"
     
-    def respond(query, history):
-        if not query.strip():
-            return history, ""
-        
-        # Add user message
-        history.append((query, None))
-        
-        # Get response
-        response = rag_system.generate_response(query)
-        
-        # Update history
-        history[-1] = (query, response)
-        return history, ""
-    
-    # Connect components
-    load_btn.click(
-        load_dataset,
-        inputs=dataset_url,
-        outputs=status
-    )
-    
-    submit_btn.click(
-        respond,
-        inputs=[query, chat_history],
-        outputs=[chatbot, query]
-    )
-    
-    query.submit(
-        respond,
-        inputs=[query, chat_history],
-        outputs=[chatbot, query]
-    )
+    def respond(message, chat_history):
+        response = rag_system.generate_response(message)
+        chat_history.append((message, response))
+        return "", chat_history
     
-    clear_btn.click(
-        lambda: [],
-        inputs=None,
-        outputs=chatbot
-    )
+    load_btn.click(load_dataset, outputs=status)
+    submit_btn.click(respond, [query, chatbot], [query, chatbot])
+    query.submit(respond, [query, chatbot], [query, chatbot])
 
 if __name__ == "__main__":
-    app.launch(debug=True)
+    app.launch(share=True)