Update app.py

app.py

@@ -44,10 +44,19 @@ class LocalLLMHandler:
                 torch.cuda.empty_cache()
                 gc.collect()
             
+            # Quantization for faster inference
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.bfloat16
+            )
+            
             model_kwargs = {
-                "device_map": "cpu",
+                "device_map": "auto", # Use GPU if available, otherwise CPU
                 "torch_dtype": torch.bfloat16,  
                 "low_cpu_mem_usage": True,
+                "quantization_config": quantization_config 
             }
             
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -63,7 +72,7 @@ class LocalLLMHandler:
 
     def generate_response(self, prompt, max_length=500):
         try:
-            inputs = self.tokenizer(prompt, return_tensors="pt")
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device) # Move inputs to the same device as the model
             outputs = self.model.generate(
                 inputs["input_ids"],
                 max_length=max_length,
@@ -78,6 +87,7 @@ class LocalLLMHandler:
             logger.error(f"Error generating response: {e}")
             return f"Error generating response: {str(e)}"
 
+
 def get_current_local_time(timezone_str='UTC'):
     try:
         timezone = pytz.timezone(timezone_str)