translator_api / main.py
jxtan's picture
Initial Commit to Test Runpod
76e433a
raw
history blame
978 Bytes
# 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])
import json
return json.dumps({"input_text": input_text, "src_code": source_language_code, "tgt_code": target_language_code})
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})