Update app.py

app.py

@@ -1,19 +1,40 @@
-from fastapi import FastAPI
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import os
 
-app = FastAPI()
+MODEL_ID = "google/gemma-3-270m-it"
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = None
+model = None
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+def load_model():
+    global tokenizer, model
+    if tokenizer is None or model is None:
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID,
+            token=HF_TOKEN
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            token=HF_TOKEN,
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True
+        )
 
-@app.post("/chat")
-def chat(prompt: str):
+def chat(prompt):
+    load_model()
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(
         **inputs,
-        max_new_tokens=200
+        max_new_tokens=128
     )
-    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"reply": reply}
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+gr.Interface(
+    fn=chat,
+    inputs="textbox",
+    outputs="textbox",
+    title="Gemma3 270M Cloud"
+).launch(server_name="0.0.0.0")