import gradio as gr
from transformers import pipeline
import json

# Load the phrase mapping from the JSON file
with open('phrase_mappings.json', 'r') as f:
    phrase_mapping = json.load(f)

# Define the UnifiedTranslator class
class UnifiedTranslator:
    def __init__(self, model_name, phrase_mapping):
        self.model = pipeline("translation", model=model_name)
        self.phrase_mapping = phrase_mapping

    def translate(self, text):
        # Check if the text is in the phrase_mapping dictionary
        if text in self.phrase_mapping:
            return self.phrase_mapping[text]
        else:
            # Use the model to translate if not in the dictionary
            return self.model(text)[0]['translation_text']

# Initialize the UnifiedTranslator with your model and custom phrases
translator = UnifiedTranslator("Bildad/Swahili-English_Translation", phrase_mapping)

# Define the Gradio interface
def translate_text(text):
    return translator.translate(text)

iface = gr.Interface(
    fn=translate_text,
    inputs="text",
    outputs="text",
    title="Swahili-English Translation",
    description="Translate Swahili to English with custom phrase mappings."
)

# Launch the interface
iface.launch(share=True)