app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import json
import os

# Ensure Hugging Face API token is set
os.environ["HF_TOKEN"] = os.getenv('hf')

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

# Mapping for language codes
lang_codes = {
    "English to Iban": ("eng_Latn", "iba_Latn"),
    "Iban to English": ("iba_Latn", "eng_Latn")
}

# Translation function
def translate_text(text, translation_direction):
    if translation_direction not in lang_codes:
        return json.dumps({"error": "Invalid translation direction"})

    src_lang, tgt_lang = lang_codes[translation_direction]

    # Tokenize input text
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)

    # Generate translation
    outputs = model.generate(**inputs, forced_bos_token_id=tokenizer.convert_tokens_to_ids(f"[{tgt_lang}]"))
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Return JSON response
    return json.dumps({
        "input_text": text,
        "translated_text": translated_text,
        "language": f"{src_lang} to {tgt_lang}"
    }, indent=4)

# Define the Gradio interface
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 and enter your text."
)

# Launch the app with a shareable link
iface.launch(share=True)