Add logging of outputs

app.py

@@ -1,28 +1,32 @@
 import gradio as gr
 from transformers import pipeline
-
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import warnings
+import weave
+
+weave.init("Faseeh")
 
 warnings.filterwarnings("ignore")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("Abdulmohsena/Faseeh")
-model = AutoModelForSeq2SeqLM.from_pretrained("Abdulmohsena/Faseeh", revision="03226648afb02523acacfe8af1363cc5ddd9e795").to(device)
+model = AutoModelForSeq2SeqLM.from_pretrained("Abdulmohsena/Faseeh").to(device)
+
+class Faseeh(weave.Model):
 
-def translate(text, temperature=0.1, tries=1):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=256).to(device)
-    outputs = model.generate(
-        **inputs,
-        do_sample=True,
-        temperature=temperature,
-        num_return_sequences=tries,
-    )
-    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    @weave.op()
+    def translate(self, input_data: str, temperature: float = 0.1):
+        inputs = tokenizer(input_data, return_tensors='pt')
+        outputs = model.generate(**inputs,
+                                 temperature=temperature,
+                                 do_sample=True)[0]
+        prediction = tokenizer.decode(outputs, skip_special_tokens=True)
+            
+        return prediction
 
-    return translation
+faseeh = Faseeh()
 
 # Gradio interface
 with gr.Blocks() as demo:
@@ -38,7 +42,7 @@ with gr.Blocks() as demo:
     translate_btn = gr.Button("Translate")
 
     # Button action
-    translate_btn.click(translate, inputs=input_text, outputs=output_text)
+    translate_btn.click(faseeh.translate, inputs=input_text, outputs=output_text)
 
 # Launch the Gradio app
 if __name__ == "__main__":