berndf/3D-text-embedding

1import streamlit as st2from transformers import AutoTokenizer, AutoModel3from sklearn.decomposition import PCA4import torch5import numpy as np6import plotly.graph_objects as go7import math8 9# Load transformer model + tokenizer..10@st.cache_resource11def load_model():12    tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")13    model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")14    return tokenizer, model15 16tokenizer, model = load_model()17 18# Encode using mean pooling19def encode_texts(texts):20    with torch.no_grad():21        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")22        output = model(**inputs)23        mask = inputs["attention_mask"].unsqueeze(-1).expand(output.last_hidden_state.shape).float()24        pooled = torch.sum(output.last_hidden_state * mask, dim=1) / mask.sum(dim=1)25        return pooled.cpu().numpy()26 27# Session state init28if "submitted_text" not in st.session_state:29    st.session_state.submitted_text = """BMW30Porsche31Mercedes32Coffee33Tea34Water35Germany36Italy37Brazil38Violin39Drums40Trumpet41Man42Women43Child"""44 45# UI layout46col1, col2 = st.columns([1, 3])47 48with col1:49    st.title("🧠 Embedding Input")50 51    with st.form(key="embedding_input_form"):52        st.form_submit_button("✅ Submit Text")53        st.text_area(54            label="Enter words (one per line)",55            key="submitted_text",56            height=400,57        )58 59texts = [t.strip() for t in st.session_state.submitted_text.split("\n") if t.strip()]60if len(texts) < 3:61    st.warning("Please enter at least three words.")62    st.stop()63 64embeddings = encode_texts(texts)65coords = PCA(n_components=3).fit_transform(embeddings)66 67# Rotation frames68frames = []69for angle in range(0, 360, 2):70    rad = math.radians(angle)71    camera = dict(eye=dict(x=2 * math.cos(rad), y=2 * math.sin(rad), z=0.7))72    frames.append(go.Frame(layout=dict(scene_camera=camera)))73 74# Plotly figure with animation controls75fig = go.Figure(76    data=[77        go.Scatter3d(78            x=coords[:, 0],79            y=coords[:, 1],80            z=coords[:, 2],81            mode="markers+text",82            text=texts,83            textposition="top center",84            textfont=dict(color="black"),85            marker=dict(size=6),86        )87    ],88    layout=go.Layout(89        title="3D Embedding Projection",90        scene=dict(91            xaxis=dict(title="X", showbackground=True, backgroundcolor="rgba(255,0,0,0.4)"),92            yaxis=dict(title="Y", showbackground=True, backgroundcolor="rgba(0,255,0,0.4)"),93            zaxis=dict(title="Z", showbackground=True, backgroundcolor="rgba(0,0,255,0.4)"),94        ),95        updatemenus=[96            dict(97                type="buttons",98                showactive=False,99                buttons=[100                    dict(101                        label="🔄 Rotate",102                        method="animate",103                        args=[104                            None,105                            dict(106                                frame=dict(duration=50, redraw=True),107                                transition=dict(duration=0),108                                fromcurrent=True,109                                mode="immediate"110                            )111                        ],112                    )113                ],114                x=0.05,115                y=0.9116            )117        ],118        margin=dict(l=0, r=0, b=0, t=30),119    ),120    frames=frames121)122 123with col2:124    st.title("📊 3D Plot")125    st.plotly_chart(fig, use_container_width=True)