Update main.py

main.py

@@ -1,22 +1,38 @@
-from fastapi import FastAPI
+from flask import Flask, request, jsonify
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-app = FastAPI()
+app = Flask(__name__)
 
 # Load model and tokenizer
 model_name = "mistralai/Mistral-7B-v0.1"
+
+# Enable quantization for better performance on free-tier Spaces
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    torch_dtype=torch.float16, 
+    device_map="auto"
+)
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-# model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-# if these 3 lines didn't work use the one above
-from transformers import BitsAndBytesConfig
-quant_config = BitsAndBytesConfig(load_in_8bit=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=quant_config, device_map="auto")
+@app.route("/")
+def home():
+    return request.url
+
+@app.route("/generate")
+def generate_text():
+    data = request.get_json()
+    prompt = data.get("prompt", "")
 
+    if not prompt:
+        return jsonify({"error": "No prompt provided"}), 400
 
-@app.post("/generate")
-async def generate_text(prompt: str):
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     outputs = model.generate(**inputs, max_length=200)
-    return {"response": tokenizer.decode(outputs[0], skip_special_tokens=True)}
+    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return jsonify({"response": response_text})
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)