Update app.py

app.py

@@ -1,17 +1,13 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoModel, AutoTokenizer, pipeline
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-# 1️⃣ Initialize app
-app = FastAPI(title="FunFox Model API")
-
-# 2️⃣ Load model and tokenizer
+# 1️⃣ Load model and tokenizer
 model_name = "DSDUDEd/funfox"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name, torch_dtype="auto")
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
 
-# 3️⃣ Setup text-generation pipeline (if applicable)
+# 2️⃣ Setup text generation pipeline
 generator = pipeline(
     task="text-generation",
     model=model,
@@ -19,18 +15,22 @@ generator = pipeline(
     device=0 if torch.cuda.is_available() else -1
 )
 
-# 4️⃣ Request schema
-class Prompt(BaseModel):
-    text: str
-    max_tokens: int = 50
+# 3️⃣ Define function for Gradio interface
+def generate_text(prompt, max_tokens=50):
+    output = generator(prompt, max_new_tokens=max_tokens)
+    return output[0]["generated_text"]
 
-# 5️⃣ API route
-@app.post("/generate")
-def generate(prompt: Prompt):
-    outputs = generator(prompt.text, max_new_tokens=prompt.max_tokens)
-    return {"generated_text": outputs[0]["generated_text"]}
+# 4️⃣ Create Gradio interface
+iface = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(label="Enter Prompt", lines=2, placeholder="Type something..."),
+        gr.Slider(minimum=10, maximum=200, step=10, label="Max Tokens")
+    ],
+    outputs=gr.Textbox(label="Generated Text"),
+    title="FunFox Text Generation",
+    description="Enter a prompt and generate text with the FunFox model."
+)
 
-# 6️⃣ Root
-@app.get("/")
-def root():
-    return {"message": "FunFox model API is running!"}
+# 5️⃣ Launch the interface
+iface.launch(share=True)