Update app.py

app.py

@@ -1,38 +1,52 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftModel
 import torch
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+
+# Setup
+app = Flask(__name__)
+CORS(app)  # Enable CORS for frontend/backend calls
+
+# Load base model + adapter
+base_model_name = "unsloth/gemma-3-12b-it-unsloth-bnb-4bit"
+adapter_name = "adarsh3601/my_gemma3_pt"
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4"
+)
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    quantization_config=bnb_config,
+    device_map="auto"
+)
+
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+model = PeftModel.from_pretrained(base_model, adapter_name)
+
+@app.route("/chat", methods=["POST"])
+def chat():
+    try:
+        data = request.json
+        prompt = data.get("message", "")
+
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=150, do_sample=True)
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        return jsonify({"response": response})
+
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+# For Hugging Face Spaces to detect the server
+@app.route("/", methods=["GET"])
+def root():
+    return "HF Space backend running"
 
-# Load the tokenizer from the Hugging Face Hub
-tokenizer = AutoTokenizer.from_pretrained("adarsh3601/my_gemma3_pt")
-
-# Load the model from Hugging Face Hub (Assuming you are using a transformer model here)
-model = AutoModelForCausalLM.from_pretrained("adarsh3601/my_gemma3_pt")
-
-# Function to generate response using the model
-def generate_response(input_text):
-    # Tokenize the input text
-    inputs = tokenizer(input_text, return_tensors="pt")
-    
-    # Generate output using the model
-    with torch.no_grad():  # Disable gradients for inference
-        outputs = model.generate(inputs['input_ids'], max_length=50)  # You can adjust max_length and other parameters
-    
-    # Decode the output and return it
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Create a Gradio interface
-def create_gradio_interface():
-    # Interface with a text input and a text output
-    interface = gr.Interface(
-        fn=generate_response,  # Function to call for generation
-        inputs=gr.Textbox(label="Enter Input Text"),  # Textbox for user input
-        outputs=gr.Textbox(label="Generated Response"),  # Textbox for output text
-        title="Text Generation with My Model",  # Title for the interface
-        description="Enter some text to generate a response using the trained model."  # Description
-    )
-    return interface
-
-# Launch the Gradio interface
 if __name__ == "__main__":
-    interface = create_gradio_interface()
-    interface.launch()
+    app.run(host="0.0.0.0", port=7860)