from translator import translator
from lang_list import LANGUAGE_NAME_TO_CODE
import runpod

def run_t2tt(input_text: str, source_language: str, target_language: str) -> str:
    source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
    target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
    out_texts, _ = translator.predict(
        input=input_text,
        task_str="T2TT",
        src_lang=source_language_code,
        tgt_lang=target_language_code,
    )
    return str(out_texts[0])

def runpod_handler(job):
    job_input = job['input']
    input_text = job_input["input_text"]
    source_language = job_input["source_language"]
    target_language = job_input["target_language"]
    return run_t2tt(input_text, source_language, target_language)

runpod.serverless.start({"handler": runpod_handler})