Update app.py

app.py

@@ -1,53 +1,36 @@
-from fastapi import FastAPI, Query
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-app = FastAPI()
+MODEL_ID = "Jackrong/Qwen3.5-9B-GLM5.1-Distill-v1"
 
-MODEL_ID = "google/gemma-4-E2B-it"
-
-print("Loading Gemma 4 on CPU...")
-
-device = torch.device("cpu")
-
-# ✅ Use tokenizer (NOT AutoProcessor)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# ✅ Load model safely on CPU
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
+    device_map="cpu",
     torch_dtype=torch.float32,
     low_cpu_mem_usage=True
-).to(device)
-
-model.eval()
-
-print("Model loaded!")
-
-@app.get("/")
-def root():
-    return {"message": "Gemma 4 API running on CPU"}
+)
 
-@app.get("/generate")
-def generate(input: str = Query(...)):
-    try:
-        inputs = tokenizer(input, return_tensors="pt").to(device)
+def chat(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt")
 
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=128,   # keep small for CPU
-                do_sample=True,
-                temperature=0.7
-            )
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=64,   # keep LOW
+        do_sample=True,
+        temperature=0.7
+    )
 
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        return {
-            "input": input,
-            "response": response
-        }
+app = gr.Interface(
+    fn=chat,
+    inputs="text",
+    outputs="text",
+    api_name="generate"
+)
 
-    except Exception as e:
-        return {"error": str(e)}
-        
+app.queue()
+app.launch()