Update app.py

app.py

@@ -23,7 +23,7 @@ def load_pipeline():
     tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left", trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         model_name, 
-        torch_dtype=torch.bfloat16,  # Use bfloat16 to reduce memory usage
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Use float16 for GPU, float32 for CPU
         device_map="auto",          # Automatically map model to available devices (e.g., GPU if available)
         trust_remote_code=True
     )
@@ -33,7 +33,7 @@ def load_pipeline():
         task="text-generation", 
         model=model, 
         tokenizer=tokenizer, 
-        torch_dtype=torch.bfloat16,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
         device_map="auto", 
         return_full_text=True  # Required for LangChain compatibility
     )