import gradio as gr
import pandas as pd
from sentence_transformers import SentenceTransformer

# 1) Load a small, fast, well-known model
model = SentenceTransformer("all-MiniLM-L6-v2")

def clean_and_embed(text: str):
    # 2) Clean: keep ASCII only & lowercase
    clean = text.encode("ascii", "ignore").decode().lower()

    # 3) Tokenize via the model’s tokenizer (quick peek at tokenization)
    tokens = model.tokenizer.tokenize(clean)

    # 4) Get the sentence embedding (as numpy array)
    emb = model.encode(clean, convert_to_numpy=True)

    # 5) Build a DataFrame: one row (the sentence) × embedding dims
    df = pd.DataFrame(
        [emb],
        index=["sentence_embedding"],
        columns=[f"dim_{i}" for i in range(emb.shape[0])]
    )
    return " ".join(tokens), df

# 6) Gradio interface
iface = gr.Interface(
    fn=clean_and_embed,
    inputs=gr.Textbox(lines=2, placeholder="Type your text here…"),
    outputs=[
        gr.Textbox(label="Tokens"),
        gr.Dataframe(label="Sentence Embedding Vector")
    ],
    title="ASCII-Clean + SentenceTransformer",
    description="Cleans input, tokenizes with a SentenceTransformer tokenizer, and shows the sentence embedding."
)

if __name__ == "__main__":
    # Important for running inside Docker (expose to the host)
    iface.launch(server_name="0.0.0.0", server_port=7860)