import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import json
import os
os.environ["HF_TOKEN"] = os.getenv('hf')

# Load the tokenizer and model from Hugging Face (authentication handled via token)
tokenizer = AutoTokenizer.from_pretrained("Darwin29/lumina-translate")
model = AutoModelForSeq2SeqLM.from_pretrained("Darwin29/lumina-translate")

# Translation function (English to Iban and Iban to English)
def translate_text(text, translation_direction):
    if translation_direction == "English to Iban":
        # Translate English to Iban
        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
        outputs = model.generate(**inputs)
        translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return json.dumps({"input_text": text, "translated_text": translated_text, "language": "en-to-iban"})
    
    elif translation_direction == "Iban to English":
        # Translate Iban to English (you'll need a reverse model or translation setup here)
        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
        outputs = model.generate(**inputs)
        translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return json.dumps({"input_text": text, "translated_text": translated_text, "language": "iban-to-en"})

# Define the Gradio interface with a dropdown to select translation direction
iface = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(label="Enter Text (English or Iban)"),
        gr.Radio(["English to Iban", "Iban to English"], label="Select Translation Direction")
    ],
    outputs=gr.JSON(),
    title="English-Iban Translator",
    description="This app translates between English and Iban. Choose the translation direction."
)

# Launch the app
iface.launch()