import gradio as gr
from transformers import AutoModel, AutoTokenizer

# Load the BAGEL model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("ByteDance-Seed/BAGEL-7B-MoT")
model = AutoModel.from_pretrained("ByteDance-Seed/BAGEL-7B-MoT")

# Define a function to generate outputs
def generate(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Create a Gradio interface
demo = gr.Interface(
    fn=generate,  # The function to run
    inputs="text",  # Input type (text box)
    outputs="text"  # Output type (text box)
)

# Launch the app
demo.launch()