Update app.py

app.py

@@ -1,7 +1,24 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+tokenizer = AutoTokenizer.from_pretrained("google/mt5-small")
+model = AutoModelForSeq2SeqLM.from_pretrained("./models/checkpoint-15000/")
+
 
 def text_processing(text):
-    return text * 2
+    inputs = [text]
+
+    # Tokenize and prepare the inputs for model
+    input_ids = tokenizer(inputs, return_tensors="pt", max_length=512, truncation=True, padding="max_length").input_ids.to(device)
+    attention_mask = tokenizer(inputs, return_tensors="pt", max_length=512, truncation=True, padding="max_length").attention_mask.to(device)
+
+    # Generate prediction
+    output = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=512)
+
+    # Decode the prediction
+    decoded_output = [tokenizer.decode(ids, skip_special_tokens=True) for ids in output]
+
+    return decoded_output[0]
 
 iface = gr.Interface(fn = text_processing, inputs='text', outputs=['text'], title='test', description='test space')