import gradio as gr
import pandas as pd
from transformers import BertTokenizer, BertModel
import torch

# Load pre-trained BERT tokenizer and model
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")

def clean_and_transform(text: str):
    # 1. Clean: remove non-ASCII & lowercase
    clean = text.encode("ascii", "ignore").decode().lower()
    
    # 2. Tokenize input text with special tokens
    inputs = tokenizer(clean, return_tensors="pt", add_special_tokens=True)
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze())
    
    # 3. Pass through BERT model
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 4. Extract embeddings for each token
    embeddings = outputs.last_hidden_state.squeeze(0).numpy()
    
    # 5. Convert embeddings into DataFrame
    df = pd.DataFrame(
        embeddings,
        index=tokens,
        columns=[f"dim_{i}" for i in range(embeddings.shape[1])]
    )
    
    return " ".join(tokens), df

# Gradio Interface
iface = gr.Interface(
    fn=clean_and_transform,
    inputs=gr.Textbox(lines=3, placeholder="Enter your text here..."),
    outputs=[
        gr.Textbox(label="Tokens"),
        gr.Dataframe(label="Embeddings", wrap=True)
    ],
    title="BERT Tokenizer & Embeddings",
    description="Enter text to see tokens and embeddings generated by BERT."
)

if __name__ == "__main__":
    iface.launch()