app.py

import time
import gradio as gr
import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

# Load model and tokenizer
print("Loading model and tokenizer...")
model_name = "panagoa/xlm-roberta-base-kbd"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
print(f"Model loaded on {device}")

def fill_mask(text, top_k=10, show_timing=False):
    """
    Fill masked tokens in the text with predictions from the model

    Args:
        text (str): Text with mask tokens
        top_k (int): Number of top predictions to return
        show_timing (bool): Whether to show inference timing

    Returns:
        str: Text with predictions for masked tokens and optional timing info
    """
    # Check if text contains mask token
    mask_token = tokenizer.mask_token
    if not text or mask_token not in text:
        return f"Please include at least one {mask_token} token in your text."

    # Measure inference time
    start_time = time.time()

    # Tokenize input
    inputs = tokenizer(text, return_tensors="pt").to(device)

    # Get token indices for mask tokens
    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]

    # Forward pass
    with torch.no_grad():
        outputs = model(**inputs)

    # Get predictions
    logits = outputs.logits

    result = []
    result.append(f"Input: {text}\n")
    result.append(f"Top {top_k} predictions for each {tokenizer.mask_token}:")

    # Process each mask token
    for i, idx in enumerate(mask_token_index):
        mask_num = i + 1
        logits_for_mask = logits[0, idx, :]

        # Get top k predictions
        top_k_tokens = torch.topk(logits_for_mask, top_k, dim=0)
        top_k_token_ids = top_k_tokens.indices.tolist()
        top_k_scores = top_k_tokens.values.tolist()

        # Convert token IDs to tokens
        top_k_tokens_text = [tokenizer.decode([token_id]) for token_id in top_k_token_ids]

        # Add to result
        result.append(f"\nMask {mask_num}:")
        for j, (token_text, score) in enumerate(zip(top_k_tokens_text, top_k_scores)):
            # Clean up token text (remove spaces and special chars)
            token_text = token_text.strip()
            result.append(f"{j+1}. {token_text} (score: {score:.4f})")

    # Calculate inference time
    inference_time = time.time() - start_time

    # Add timing information if requested
    if show_timing:
        result.append(f"\n[Inference time: {inference_time:.3f} seconds]")

    return "\n".join(result)

# Example sentences with masks - use the tokenizer's mask token
mask_token = tokenizer.mask_token
EXAMPLES = [
    f"Iэджи {mask_token} абы пэжу.",
    f"Абы арэзы {mask_token} хъунукъым.",
    f"Си гум {mask_token} ахэр.",
    f"Сэ {mask_token} уэ уи фэм дэкIар.",
    f"НтIэ, арати, зы {mask_token} зы фызыжьрэ псэурт.",
    f"Усэ зыбжани {mask_token} адыгэ псори зыгъэгузавэ а темэм.",
    f"УщIалэщ, дяпэкIэщ уи насып {mask_token}.",
    f"Уэ езым уи нитIкIэ {mask_token}, къэпцIыхунщ.",
    f"Уэ уи лъэпкъым иджыри куэд {mask_token}.",
]

# Create Gradio interface
with gr.Blocks(title="XLM-RoBERTa for Kabardian Language") as demo:
    gr.Markdown("# 🔍 Kabardian Word Prediction")
    gr.Markdown(f"This app uses the [panagoa/xlm-roberta-base-kbd](https://huggingface.co/panagoa/xlm-roberta-base-kbd) model to predict missing words in Kabardian text. Add the special {mask_token} token where you want the model to predict words.")

    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label=f"Enter text with {mask_token} token",
                placeholder=f"Адыгэр зи {mask_token} лъэпкъщ.",
                lines=3
            )

            top_k = gr.Slider(
                minimum=5,
                maximum=20,
                value=5,
                step=5,
                label="Number of predictions"
            )

            show_timing = gr.Checkbox(
                label="Show Inference Time",
                value=False
            )

            predict_btn = gr.Button("Predict", variant="primary")

        with gr.Column():
            text_output = gr.Textbox(
                label="Prediction Results",
                lines=15
            )

    gr.Examples(
        EXAMPLES,
        inputs=[text_input],
        label=f"Example Kabardian sentences with {mask_token}"
    )

    # Set up button actions
    predict_btn.click(
        fn=fill_mask,
        inputs=[text_input, top_k, show_timing],
        outputs=text_output
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()