feat: Add experimental version of plot after adaption from referred space

app.py

@@ -1,13 +1,75 @@
+from typing import List
+
+import numpy as np
+
 import streamlit as st
 import tweepy
-from sentence_transformers import SentenceTransformer
 
+from bokeh.models import ColumnDataSource, HoverTool
+from bokeh.palettes import Cividis256 as Pallete
+from bokeh.plotting import Figure, figure
+from bokeh.transform import factor_cmap
+
+from sklearn.manifold import TSNE
+from sentence_transformers import SentenceTransformer
 
+# Original implementation from: https://huggingface.co/spaces/edugp/embedding-lenses/blob/main/app.py
 @st.cache(show_spinner=False, allow_output_mutation=True)
 def load_model(model_name: str) -> SentenceTransformer:
     embedder = model_name
     return SentenceTransformer(embedder)
 
+def embed_text(text: List[str], model: SentenceTransformer) -> np.ndarray:
+    return model.encode(text)
+
+def get_tsne_embeddings(
+    embeddings: np.ndarray, perplexity: int = 30, n_components: int = 2, init: str = "pca", n_iter: int = 5000, random_state: int = SEED
+) -> np.ndarray:
+    tsne = TSNE(perplexity=perplexity, n_components=n_components, init=init, n_iter=n_iter, random_state=random_state)
+    return tsne.fit_transform(embeddings)
+
+def draw_interactive_scatter_plot(
+    texts: np.ndarray, xs: np.ndarray, ys: np.ndarray, values: np.ndarray, labels: np.ndarray, text_column: str, label_column: str
+) -> Figure:
+    # Normalize values to range between 0-255, to assign a color for each value
+    max_value = values.max()
+    min_value = values.min()
+    if max_value - min_value == 0:
+        values_color = np.ones(len(values))
+    else:
+        values_color = ((values - min_value) / (max_value - min_value) * 255).round().astype(int).astype(str)
+    values_color_set = sorted(values_color)
+    values_list = values.astype(str).tolist()
+    values_set = sorted(values_list)
+    labels_list = labels.astype(str).tolist()
+    source = ColumnDataSource(data=dict(x=xs, y=ys, text=texts, label=values_list, original_label=labels_list))
+    hover = HoverTool(tooltips=[(text_column, "@text{safe}"), (label_column, "@original_label")])
+    p = figure(plot_width=800, plot_height=800, tools=[hover])
+    p.circle("x", "y", size=10, source=source, fill_color=factor_cmap("label", palette=[Pallete[int(id_)] for id_ in values_color_set], factors=values_set))
+    p.axis.visible = False
+    p.xgrid.grid_line_color = None
+    p.ygrid.grid_line_color = None
+    p.toolbar.logo = None
+    return p
+
+# Up to here
+def generate_plot(
+    df: List[str],
+    labels: List[int],
+    model: SentenceTransformer,
+) -> Figure:
+    with st.spinner(text="Embedding text..."):
+        embeddings = embed_text(df, model)
+    logger.info("Encoding labels")
+    encoded_labels = encode_labels(labels)
+    with st.spinner("Reducing dimensionality..."):
+        embeddings_2d = get_tsne_embeddings(embeddings)
+    logger.info("Generating figure")
+    plot = draw_interactive_scatter_plot(
+        df, embeddings_2d[:, 0], embeddings_2d[:, 1], encoded_labels.values, labels, 'text', 'label'
+    )
+    return plot
+
 
 client = tweepy.Client(bearer_token=st.secrets["tw_bearer_token"])
 model_to_use = {
@@ -21,14 +83,35 @@ col1, col2 = st.columns(2)
 with col1:
     tw_user = st.text_input("Twitter handle", "huggingface")
 with col2:
-    sample = st.number_input("Maximum number of tweets to use", 1, 300, 100, 10)
+    tw_sample = st.number_input("Maximum number of tweets to use", 1, 300, 100, 10)
     
 expected_lang = st.radio(
     "What language should be assumed to be found?",
-    ('English', 'Use all the ones you know (~15 lang)'))
+    ('English', 'Use all the ones you know (~15 lang)'),
+    0
+)
 
+with st.spinner(text="Loading model..."):
+    model = load_model(model_to_use[expected_lang])
 
 usr = client.get_user(username=tw_user)
  
-st.write(usr.data.id)
- 
+# st.write(usr.data.id)
+
+if tw_user:
+    with st.spinner(f"Getting to know the '{tw_user}'..."):
+        tweets_objs = []
+        while tw_sample >= 100:
+            current_sample = min(100, tw_sample)
+            tweets_response = client.get_user_tweets(usr.data.id, max_results=current_sample)
+            tweets_objs += tweets_response.data
+            tw_sample -= current_sample
+        tweets_response = client.get_user_tweets(usr.data.id, max_results=tw_sample)
+        tweets_objs += tweets_response.data
+    tweets_txt = [tweet.text for tweet in tweets_objs]
+    labels = [0] * len(tweets_txt)
+    # plot = generate_plot(df, text_column, label_column, sample, dimensionality_reduction_function, model)
+    plot = generate_plot(tweets_txt, labels, model)
+    logger.info("Displaying plot")
+    st.bokeh_chart(plot)
+    logger.info("Done")