import torch
import gradio as gr
import json
from transformers import pipeline

# Use a pipeline with the specified model for translation
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, source_language, destination_language):
    src_code = get_FLORES_code_from_language(source_language)
    dest_code = get_FLORES_code_from_language(destination_language)
    translation = text_translator(text, src_lang=src_code, tgt_lang=dest_code)
    return translation[0]["translation_text"]

gr.close_all()

# Define the language options based on the JSON file
language_options = [entry['Language'] for entry in language_data]

demo = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(label="Input text to translate", lines=6),
        gr.Dropdown(language_options, label="Select Source Language"),
        gr.Dropdown(language_options, label="Select Destination Language")
    ],
    outputs=[gr.Textbox(label="Translated text", lines=4)],
    title="Multi-language Translator",
    description="THIS APPLICATION WILL BE USED TO TRANSLATE TEXT BETWEEN MULTIPLE LANGUAGES."
)

demo.launch()