Update app.py

app.py

@@ -1,3 +1,5 @@
+import gradio as gr
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 import torch
 from fastapi import FastAPI
 from pydantic import BaseModel
@@ -5,17 +7,16 @@ from threading import Thread
 import uvicorn
 
 # Initialize FastAPI
-
 app = FastAPI()
 
 # Load the tokenizer and model
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
+tokenizer = T5Tokenizer.from_pretrained("alpeshsonar/lot-t5-small-filter", legacy=False)
+model = T5ForConditionalGeneration.from_pretrained("alpeshsonar/lot-t5-small-filter")
 
 # Define the function to generate text for Gradio
 def generate_text(input_text):
     input_text = "Extract lots from given text.\n" + input_text
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    inputs = tokenizer.encode(input_text, return_tensors="pt")
     outputs = model.generate(inputs, max_new_tokens=1024)
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return result
@@ -30,10 +31,6 @@ iface = gr.Interface(
 )
 
 # Define a request body model for FastAPI
-
-
-
-
 class TextInput(BaseModel):
     input_text: str
 
@@ -41,12 +38,9 @@ class TextInput(BaseModel):
 @app.post("/generate")
 async def generate_text_api(input_data: TextInput):
     input_text = input_data.input_text
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
-
-
+    inputs = tokenizer.encode(input_text, return_tensors="pt")
     outputs = model.generate(inputs, max_new_tokens=1024)
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
     return {"output": result}
 
 # Health check endpoint