import gradio as gr
from transformers import AutoTokenizer, BartModel
import torch

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
model = BartModel.from_pretrained("facebook/bart-base")

def generate_response(text):
    # Tokenize the input text
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)

    # Generate model output
    outputs = model(**inputs)

    # Get the last hidden states from the model output
    last_hidden_states = outputs.last_hidden_state

    # Convert tensor to a human-readable format
    output_text = "\n".join([str(token) for token in last_hidden_states[0][0]])

    # Return last hidden states as output
    return output_text

iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=7, label="Input Text"),
    outputs=gr.Textbox(label="Output Last Hidden States")
)

iface.launch()