app.py

import gradio as gr
from comet import download_model, load_from_checkpoint
import os

model_path = os.environ.get("HF_MODEL_PATH", download_model("wasanx/ComeTH"))
model = load_from_checkpoint(model_path)

def score_translation(src_text, mt_text):
    translations = [{"src": src_text, "mt": mt_text}]
    results = model.predict(translations, batch_size=1, gpus=1)
    return results['scores'][0]

examples = [
    ["The weather is beautiful today.", "วันนี้อากาศดีมาก"],
    ["I need to go to the hospital.", "ฉันต้องไปโรงพยาบาล"],
    ["This restaurant serves delicious food.", "ร้านอาหารนี้เสิร์ฟอาหารอร่อย"],
    ["Can you help me find the nearest train station?", "คุณช่วยฉันหาสถานีรถไฟที่ใกล้ที่สุดได้ไหม"]
]

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# ComeTH Translation Quality Evaluator")
    
    with gr.Row():
        with gr.Column():
            src_input = gr.Textbox(label="Source Text (English)", placeholder="Enter English text here...")
            mt_input = gr.Textbox(label="Machine Translation (Thai)", placeholder="Enter Thai translation here...")
            score_button = gr.Button("Evaluate Translation", variant="primary")
        
        with gr.Column():
            score_output = gr.Label(label="Quality Scores")
            gr.Markdown("### Higher scores indicate better translation quality across multiple dimensions")
    
    gr.Examples(examples=examples, inputs=[src_input, mt_input], outputs=score_output, fn=score_translation)
    
    score_button.click(fn=score_translation, inputs=[src_input, mt_input], outputs=score_output)

if __name__ == "__main__":
    demo.launch()