Update to use Gradio instead of FastAPI

app.py

@@ -1,5 +1,4 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
@@ -9,14 +8,8 @@ model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5p-220m")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
-app = FastAPI()
-
-class Input(BaseModel):
-    prompt: str
-
-@app.post("/generate")
-async def generate(input: Input):
-    inputs = tokenizer(input.prompt, return_tensors="pt").to(device)
+def generate(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
     outputs = model.generate(
         **inputs,
         max_length=2048,
@@ -27,4 +20,17 @@ async def generate(input: Input):
         pad_token_id=tokenizer.pad_token_id,
     )
     output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"output": output_text}
+    return output_text
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate,
+    inputs=gr.Textbox(lines=10, label="Input Prompt"),
+    outputs=gr.Textbox(label="Generated Output"),
+    title="LLaMA 7B Server",
+    description="A web interface for interacting with the LLaMA 7B model."
+)
+
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -2,3 +2,4 @@ fastapi
 uvicorn[standard]
 transformers
 torch
+gradio>=4.17.0

start.sh

@@ -1,2 +1,3 @@
 #!/bin/bash
+python app.py
 uvicorn app:app --host 0.0.0.0 --port 7860