app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch
import os

# Login to Hugging Face Hub
access_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
login(token=access_token)


# Load model and tokenizer from the Hugging Face Hub
model_name = "kuyesu22/sunbird-ug-lang-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
model = PeftModel.from_pretrained(model, model_name)

# Ensure the model is in evaluation mode
model.eval()

# Define the translation function
def translate(text, source_lang="Runyankole", target_lang="English"):
    prompt = f"Translate from {source_lang} to {target_lang}: {text}"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    
    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            max_length=100,
            num_beams=5,
            early_stopping=True
        )
    
    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translation

# Set up the Gradio interface
def runyankole_to_english(text):
    return translate(text, source_lang="Runyankole", target_lang="English")

def english_to_runyankole(text):
    return translate(text, source_lang="English", target_lang="Runyankole")

# Create Gradio inputs and interface
with gr.Blocks() as demo:
    gr.Markdown("# Runyankole-English Translation Model")
    
    with gr.Tab("Runyankole to English"):
        runyankole_input = gr.Textbox(label="Enter Runyankole Text")
        english_output = gr.Textbox(label="English Translation")
        gr.Button("Translate").click(runyankole_to_english, inputs=runyankole_input, outputs=english_output)
    
    with gr.Tab("English to Runyankole"):
        english_input = gr.components.Textbox(label="Enter English Text")
        runyankole_output = gr.components.Textbox(label="Runyankole Translation")
        gr.Button("Translate").click(english_to_runyankole, inputs=english_input, outputs=runyankole_output)

# Launch the app
demo.launch()