Update app.py

app.py

@@ -1,34 +1,30 @@
 import gradio as gr
 from comet import download_model, load_from_checkpoint
 
-model_path = download_model("wasanx/ComeTH")
-model = load_from_checkpoint(model_path)
+def load_cometh_model():
+    model_path = download_model("wasanx/ComeTH")
+    return load_from_checkpoint(model_path)
 
-def score_translation(src_text, mt_text):
-    translations = [{
-        "src": src_text,
-        "mt": mt_text,
-    }]
+def evaluate_translation(src_text, mt_text):
+    translations = [{"src": src_text, "mt": mt_text}]
     results = model.predict(translations, batch_size=1, gpus=1)
-    scores = results.get('scores', [[]])[0]
-    output = {f"Score {i+1}": score for i, score in enumerate(scores)}
-    return output
+    return float(results['scores'][0])
 
-def main():
-    with gr.Blocks() as demo:
-        gr.Markdown("# Translation Quality Scoring with ComeTH Model")
-        with gr.Row():
-            with gr.Column():
-                src_input = gr.Textbox(label="Source Text (English)")
-                mt_input = gr.Textbox(label="Machine Translation Text (Thai)")
-            with gr.Column():
-                score_output = gr.Label(num_top_classes=5, label="Quality Scores")
-        score_button = gr.Button("Score Translation")
-        score_button.click(fn=score_translation,
-                           inputs=[src_input, mt_input],
-                           outputs=[score_output])
+model = load_cometh_model()
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo = gr.Interface(
+    fn=evaluate_translation,
+    inputs=[
+        gr.Textbox(label="English Source Text"),
+        gr.Textbox(label="Thai Translation")
+    ],
+    outputs=gr.Number(label="Quality Score"),
+    examples=[
+        ["This is a test sentence.", "นี่คือประโยคทดสอบ"],
+        ["The weather is nice today.", "อากาศดีมากวันนี้"]
+    ],
+    title="ComeTH Translator Evaluator"
+)
 
 if __name__ == "__main__":
-    main()
+    demo.launch()