import torch
import gradio as gr
import json

# Use a pipeline as a high-level helper
from transformers import pipeline

text_translator = pipeline("translation", model="facebook/nllb-200-distilled-600M",
                           torch_dtype=torch.bfloat16)

# Load the JSON data from the file
with open('language.json', 'r') as file:
    language_data = json.load(file)

def get_FLORES_code_from_language(language):
    for entry in language_data:
        if entry['Language'].lower() == language.lower():
            return entry['FLORES-200 code']
    return None


def translate_text(text, destination_language):
    # text = "Hello Friends, How are you?"
    dest_code= get_FLORES_code_from_language(destination_language)
    translation = text_translator(text,
                                  src_lang="eng_Latn",
                                  tgt_lang=dest_code)
    return translation[0]["translation_text"]

gr.close_all()

demo = gr.Interface(fn=translate_text,
                    inputs=[gr.Textbox(label="Zu übersetzenden Text eingeben",lines=6), 
                            gr.Dropdown(["German","French", "Hindi", "Romanian	"], 
                            label="Zielsprache auswählen")],
                    outputs=[gr.Textbox(label="Übersetzter Text",lines=4)],
                    title="Projekt 4: Mehrsprachiger Übersetzer",
                    description="DIESE ANWENDUNG WIRD VERWENDET, UM EINEN BELIEBIGEN ENGLISCHEN TEXT IN MEHRERE SPRACHEN ZU ÜBERSETZEN",
                    allow_flagging="never",
                    submit_btn="Übermitteln",
                    clear_btn="Bereinigen")
demo.launch()