Update app.py

app.py

@@ -1,59 +1,32 @@
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
-import torch
-from flask import Flask, request, jsonify
-from flask_cors import CORS
+import gradio as gr
 
-# Setup
-app = Flask(__name__)
-CORS(app)
-
-# Model details
-base_model_name = "unsloth/gemma-3-12b-it-unsloth-bnb-4bit"  # The model you are using
+# Model loading
+base_model_name = "unsloth/gemma-3-12b-it-unsloth-bnb-4bit"
 adapter_name = "adarsh3601/my_gemma3_pt"
-
-# 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={"": device},
-    torch_dtype=torch.float16,  # Use float16 for efficient GPU usage
-    load_in_4bit=True  # Enable 4-bit quantization for reduced memory usage
+    torch_dtype=torch.float16,
+    load_in_4bit=True
 )
 
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 model = PeftModel.from_pretrained(base_model, adapter_name)
-
-# Move model to the GPU
 model.to(device)
 
-@app.route("/chat", methods=["POST"])
-def chat():
-    try:
-        data = request.json
-        prompt = data.get("message", "")
-
-        # 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)
-
-        # Decode the output and return the response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        return jsonify({"response": response})
-
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
-
-@app.route("/", methods=["GET"])
-def root():
-    return "HF Space backend running"
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
-
+# Chat function
+def chat(message):
+    inputs = tokenizer(message, return_tensors="pt")
+    inputs = {k: v.to(device).half() for k, v in inputs.items()}
+    outputs = model.generate(**inputs, max_new_tokens=150, do_sample=True)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+# Launch Gradio app
+iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Gemma Chatbot")
+iface.launch()