import gradio as gr
import sys
import torch

try:
    from unbabel.comet import download_model, load_from_checkpoint
except ImportError:
    print("Error: unbabel-comet package not installed")
    print("Install with: pip install unbabel-comet torch gradio")
    sys.exit(1)

def evaluate_translation(src_text, mt_text):
    if not hasattr(evaluate_translation, "model"):
        try:
            model_path = download_model("wasanx/ComeTH")
            evaluate_translation.model = load_from_checkpoint(model_path)
        except Exception as e:
            return f"Error loading model: {str(e)}"
    
    translations = [{"src": src_text, "mt": mt_text}]
    results = evaluate_translation.model.predict(
        translations, 
        batch_size=1, 
        gpus=0 if not torch.cuda.is_available() else 1
    )
    return float(results['scores'][0])

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__":
    demo.launch()