import gradio as gr
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch._dynamo
torch._dynamo.config.suppress_errors = True

# Load the model and tokenizer
model_id = "answerdotai/ModernBERT-base"  # Replace with your conversational model if needed
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForMaskedLM.from_pretrained(model_id)

# Function for conversation
def conversation(input_text):
    # Prepare the input text with a [MASK] token for a masked language model
    inputs = tokenizer(input_text, return_tensors="pt")
    
    # Generate predictions
    outputs = model(**inputs)
    
    masked_index = inputs["input_ids"][0].tolist().index(tokenizer.mask_token_id)
    predicted_token_id = outputs.logits[0, masked_index].argmax(axis=-1)
    predicted_token = tokenizer.decode(predicted_token_id)

    return f"Predicted response: {predicted_token}"

# Define the Gradio interface
interface = gr.Interface(
    fn=conversation,
    inputs=gr.Textbox(label="Enter your text (include [MASK]):"),
    outputs=gr.Textbox(label="Predicted Response"),
    title="Masked Language Model Conversation",
    description="Type a sentence with [MASK] to predict the masked word using ModernBERT."
)

# Launch the interface
interface.launch()