Update app.py

app.py

@@ -3,28 +3,32 @@ from peft import PeftModel
 import torch
 from flask import Flask, request, jsonify
 from flask_cors import CORS
+from huggingface_hub import login
 
 # Setup
 app = Flask(__name__)
 CORS(app)
 
 # Model details
-base_model_name = "unsloth/gemma-3-12b-it-unsloth-bnb-4bit"
+base_model_name = "unsloth/gemma-3-12b-it-unsloth-bnb-4bit"  # The model you are using
 adapter_name = "adarsh3601/my_gemma3_pt"
 
-# Load the base model on GPU with 4-bit quantization
+# Use CUDA for GPU acceleration (Nvidia T4 small supports CUDA)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load the base model with quantization enabled for the GPU
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_name,
-    device_map="auto",  # Automatically choose GPU if available
-    load_in_4bit=True,
-    torch_dtype=torch.float16
+    device_map={"": device},
+    torch_dtype=torch.float16,  # Use float16 for efficient GPU usage
+    load_in_4bit=True  # Enable 4-bit quantization for reduced memory usage
 )
 
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-
-# Load adapter on top of base model
 model = PeftModel.from_pretrained(base_model, adapter_name)
-model.eval()
+
+# Move model to the GPU
+model.to(device)
 
 @app.route("/chat", methods=["POST"])
 def chat():
@@ -32,13 +36,14 @@ def chat():
         data = request.json
         prompt = data.get("message", "")
 
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        # Tokenize the input and move it to GPU
+        inputs = tokenizer(prompt, return_tensors="pt")
+        inputs = {k: v.to(device).half() for k, v in inputs.items()}  # Ensure inputs are in float16
+
+        # Generate the response using the model
+        outputs = model.generate(**inputs, max_new_tokens=150, do_sample=True)
 
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=150,
-            do_sample=True
-        )
+        # Decode the output and return the response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
         return jsonify({"response": response})