Update app.py

app.py

@@ -6,7 +6,8 @@ import faiss
 from sentence_transformers import SentenceTransformer
 from datasets import load_dataset
 from dotenv import load_dotenv
-import threading
+import asyncio
+import time
 
 # Load environment variables
 load_dotenv()
@@ -24,7 +25,7 @@ class GeminiRAGSystem:
         self.chunks = []
         self.dataset_loaded = False
         self.loading_error = None
-        self.gemini_api_key = os.getenv("AIzaSyASrFvE3gFPigihza0JTuALzZmBx0Kc3d0")
+        self.gemini_api_key = os.getenv("GEMINI_API_KEY")
         
         # Initialize embedding model
         try:
@@ -36,44 +37,40 @@ class GeminiRAGSystem:
         if self.gemini_api_key:
             genai.configure(api_key=self.gemini_api_key)
         
-        # Start dataset loading in background
-        self.load_dataset_in_background()
+        # Start dataset loading
+        self.load_dataset()
     
-    def load_dataset_in_background(self):
-        """Load dataset in a background thread"""
-        def load_task():
-            try:
-                # Load dataset directly
-                dataset = load_dataset(
-                    DATASET_NAME,
-                    split='train',
-                    download_mode="force_redownload"  # Fixes extraction error
-                )
-                
-                # Process dataset
-                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("Dataset must have 'text' or 'context' field")
-                
-                # Create embeddings
-                embeddings = self.embedding_model.encode(
-                    self.chunks,
-                    show_progress_bar=False,
-                    convert_to_numpy=True
-                )
-                self.index = faiss.IndexFlatL2(embeddings.shape[1])
-                self.index.add(embeddings.astype('float32'))
-                
-                self.dataset_loaded = True
-            except Exception as e:
-                self.loading_error = str(e)
-                print(f"Dataset loading failed: {str(e)}")
-        
-        # Start the loading thread
-        threading.Thread(target=load_task, daemon=True).start()
+    def load_dataset(self):
+        """Load dataset synchronously"""
+        try:
+            # Load dataset directly
+            dataset = load_dataset(
+                DATASET_NAME,
+                split='train',
+                download_mode="force_redownload"
+            )
+            
+            # Process dataset
+            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("Dataset must have 'text' or 'context' field")
+            
+            # Create embeddings
+            embeddings = self.embedding_model.encode(
+                self.chunks,
+                show_progress_bar=False,
+                convert_to_numpy=True
+            )
+            self.index = faiss.IndexFlatL2(embeddings.shape[1])
+            self.index.add(embeddings.astype('float32'))
+            
+            self.dataset_loaded = True
+        except Exception as e:
+            self.loading_error = str(e)
+            print(f"Dataset loading failed: {str(e)}")
     
     def get_relevant_context(self, query: str) -> str:
         """Retrieve most relevant chunks"""
@@ -96,10 +93,10 @@ class GeminiRAGSystem:
         """Generate response with robust error handling"""
         if not self.dataset_loaded:
             if self.loading_error:
-                return f"Dataset loading failed: {self.loading_error}"
-            return "Dataset is still loading, please wait..."
+                return f"⚠️ Dataset loading failed: {self.loading_error}"
+            return "⚠️ Dataset is still loading, please wait..."
         if not self.gemini_api_key:
-            return "Please set your Gemini API key in environment variables"
+            return "🔑 Please set your Gemini API key in environment variables"
         
         context = self.get_relevant_context(query)
         if not context:
@@ -116,7 +113,7 @@ class GeminiRAGSystem:
             response = model.generate_content(prompt)
             return response.text
         except Exception as e:
-            return f"API Error: {str(e)}"
+            return f"⚠️ API Error: {str(e)}"
 
 # Initialize system
 try:
@@ -129,31 +126,29 @@ with gr.Blocks(title="UE Chatbot") as app:
     gr.Markdown("# UE 24 Hour Service")
     
     with gr.Row():
-        chatbot = gr.Chatbot(height=500, label="Chat History", 
-                           avatar_images=(None, (None, "https://www.google.com/imgres?q=ue%20lahore%20uoe%20image%20logo&imgurl=https%3A%2F%2Fblogger.googleusercontent.com%2Fimg%2Fb%2FR29vZ2xl%2FAVvXsEglo3Aj2gTXX2j9LIG89CiFB3uj8hS4pkueQtWZrBpLkHk00zbeXuMU1soAxFWz0tenyRIwL7-YxTS9g14DSQnvgGFRGqoWMP2otVtbKbJvQ28FVpUiAzkScgA57_EQiVZ_FOSuS9cGvWg%2Fw1200-h630-p-k-no-nu%2FUE.Logo.jpg&imgrefurl=http%3A%2F%2Fuejauharabad.blogspot.com%2F2012%2F10%2Fue-logo.html&docid=ZJhky4S29RFw9M&tbnid=sopb9CZYLJPitM&vet=12ahUKEwjv7r322aqMAxXM_7sIHSTUAnYQM3oECE8QAA..i&w=769&h=403&hcb=2&ved=2ahUKEwjv7r322aqMAxXM_7sIHSTUAnYQM3oECE8QAA")),
-                           bubble_full_width=False)
+        chatbot = gr.Chatbot(
+            height=500,
+            avatar_images=(None, (None, "https://huggingface.co/spaces/groq/Groq-LLM/resolve/main/groq_logo.png")),
+            bubble_full_width=False
+        )
     
     with gr.Row():
-        query = gr.Textbox(label="Your question", 
-                         placeholder="Ask your question...", 
-                         scale=4)
+        query = gr.Textbox(
+            label="Your question",
+            placeholder="Ask your question...",
+            scale=4
+        )
         submit_btn = gr.Button("Submit", variant="primary", scale=1)
     
     with gr.Row():
         clear_btn = gr.Button("Clear Chat", variant="secondary")
 
     # Status indicator
-    status = gr.Textbox(label="System Status", 
-                       value="Initializing...",
-                       interactive=False)
-
-    # Update status periodically
-    def update_status():
-        if rag_system.loading_error:
-            return f"Error: {rag_system.loading_error}"
-        return "Ready" if rag_system.dataset_loaded else "Loading dataset..."
-    
-    app.load(update_status, None, status, every=1)
+    status = gr.Textbox(
+        label="System Status",
+        value="Loading dataset..." if not rag_system.dataset_loaded else "Ready",
+        interactive=False
+    )
 
     # Event handlers
     def respond(message, chat_history):