Update app.py

app.py

@@ -1,8 +1,15 @@
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
+import torch
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+
+# Initialize the FastAPI app
+app = FastAPI()
 
 # Load the tokenizer and model
-tokenizer = T5Tokenizer.from_pretrained("alpeshsonar/lot-t5-small-filter")
+tokenizer = T5Tokenizer.from_pretrained("alpeshsonar/lot-t5-small-filter", legacy=False)
 model = T5ForConditionalGeneration.from_pretrained("alpeshsonar/lot-t5-small-filter")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
@@ -26,3 +33,25 @@ iface = gr.Interface(
 
 # Launch the interface
 iface.launch()
+
+
+# Define a request body model
+class TextInput(BaseModel):
+    input_text: str
+
+# Endpoint to generate text from the model
+@app.post("/generate")
+async def generate_text(input_data: TextInput):
+    input_text = input_data.input_text
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+
+    # Generate the output
+    outputs = model.generate(inputs, max_new_tokens=1024)
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return {"output": result}
+
+# Health check endpoint
+@app.get("/health")
+async def health_check():
+    return {"status": "API is running"}