Update app.py

app.py

@@ -11,9 +11,13 @@ class RAGChatbot:
                  embedding_model="all-MiniLM-L6-v2"):
         # Initialize tokenizer and model
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.bnb_config = BitsAndBytesConfig(
+                        load_in_8bit=True,   # Enable 8-bit loading
+                        llm_int8_threshold=6.0,  # Threshold for mixed-precision computation
+                    )
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name, 
-            torch_dtype=torch.float16, 
+            quantization_config=bnb_config, 
             device_map="auto"
         )
         
@@ -88,12 +92,19 @@ class RAGChatbot:
             # Generate response
             response = self.generate_response(query, context)
             
-            # Append to history and return as list of tuples
-            updated_history = history + [[query, response]]
+            # Append to history using messages format
+            updated_history = history + [
+                {"role": "user", "content": query},
+                {"role": "assistant", "content": response}
+            ]
             return updated_history, ""
         except Exception as e:
             error_response = f"An error occurred: {str(e)}"
-            return history + [[query, error_response]], ""
+            updated_history = history + [
+                {"role": "user", "content": query},
+                {"role": "assistant", "content": error_response}
+            ]
+            return updated_history, ""
 
 # Create Gradio interface
 def create_interface():
@@ -108,7 +119,7 @@ def create_interface():
         
         status_output = gr.Textbox(label="Load Status")
         
-        chatbot = gr.Chatbot()
+        chatbot = gr.Chatbot(type="messages")  # Specify message type
         msg = gr.Textbox(label="Enter your query")
         submit_btn = gr.Button("Send")
         clear_btn = gr.Button("Clear Chat")