File size: 1,283 Bytes
4c33ae7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
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() |