Update app.py

app.py

@@ -1,64 +1,23 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-MODEL_NAME = "Mayur74/tinyllama-finetune-mayur"
-
-app = FastAPI()
-
-# -----------------------------
-# LOAD MODEL
-# -----------------------------
-print("Loading model...")
-
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_NAME,
-    use_auth_token=True
-)
-
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    use_auth_token=True
-)
-
-print("Model loaded successfully.")
-
-
-# -----------------------------
-# INPUT SCHEMA
-# -----------------------------
-class InputData(BaseModel):
-    prompt: str
-    max_new_tokens: int = 150
-    temperature: float = 0.7
-
-
-# -----------------------------
-# HOME ROUTE
-# -----------------------------
-@app.get("/")
-def home():
-    return {"status": "running", "message": "TinyLlama API Online 🚀"}
-
-
-# -----------------------------
-# GENERATION ENDPOINT
-# -----------------------------
-@app.post("/generate")
-def generate_text(data: InputData):
-
-    inputs = tokenizer(data.prompt, return_tensors="pt").to(model.device)
+import gradio as gr
 
+def generate(prompt, max_new_tokens=80, temperature=0.7):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     output = model.generate(
         **inputs,
-        max_new_tokens=data.max_new_tokens,
-        temperature=data.temperature,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
         do_sample=True
     )
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
+iface = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Slider(1, 500, value=80, label="Max New Tokens"),
+        gr.Slider(0, 1, value=0.7, label="Temperature")
+    ],
+    outputs="text"
+)
 
-    result = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    return {"response": result}
+iface.launch()