from gradio_client import Client
from iso639 import languages


class SalamandraTA7bTranslator:
    def __init__(self, hf_token):
        self.client = Client("BSC-LT/SalamandraTA-7B-Demo", hf_token=hf_token)

    def translate(self, text, source_lang, target_lang):
        if not text:
            return ""

        # we assume that they are specifying the language by code so we need to convert it to name
        lang1 = languages.get(alpha2=source_lang).name
        lang2 = languages.get(alpha2=target_lang).name
        result = self.client.predict(
            task="Translation",
            source=lang1,
            target=lang2,
            input_text=text,
            mt_text=None,
            api_name="/generate_output"
        )
        return result[0]