Add demo

.gitignore

@@ -0,0 +1,2 @@
+.DS_Store
+__pycache__/

app.py

@@ -0,0 +1,34 @@
+import streamlit as st
+from utils import Recommender
+import pandas as pd
+from transformers import RobertaModel
+
+st.title("Quotes RecSys")
+
+@st.cache_resource
+def get_recommender():
+    return Recommender(pd.read_csv('data/quotes.csv'),
+                "cardiffnlp/twitter-roberta-base-emotion-multilabel-latest",
+                base_model=RobertaModel, ckpt="models/twitter.pt")
+
+
+recommender = get_recommender()
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+if prompt := st.chat_input("What is up?"):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+        quote, author = recommender.recommend(prompt)
+        full_response = f"> {quote}\n\n _{author}_"
+        message_placeholder.markdown(full_response)
+    st.session_state.messages.append({"role": "assistant", "content": full_response})

requirements.txt

@@ -0,0 +1,3 @@
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch
+transformers

utils.py

@@ -0,0 +1,68 @@
+import random
+import torch
+import numpy as np
+import pandas as pd
+from stqdm import stqdm
+from torch import nn
+from torch.nn import functional as F
+from transformers import AutoTokenizer, AutoModel
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+class DSSM(nn.Module):
+    def __init__(self, base_model_name, base_model=AutoModel):
+        super().__init__()
+        self.diary_emb = base_model.from_pretrained(base_model_name, add_pooling_layer=False)
+        self.quote_emb = base_model.from_pretrained(base_model_name, add_pooling_layer=False)
+
+    def forward(self, diary, quote):
+        return self.diary_emb(**diary), self.quote_emb(**quote)
+
+
+def get_models_and_tokenizer(base_model_name, base_model=AutoModel, ckpt=None):
+    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+    model = DSSM(base_model_name, base_model=base_model)
+    if ckpt:
+        print("use ckpt")
+        model.load_state_dict(torch.load(ckpt, map_location=device))
+    model.to(device)
+    return model.diary_emb, model.quote_emb, tokenizer
+
+
+def model_inference(model, tokenizer, text):
+    tokenized_text = tokenizer(text, return_tensors="pt", truncation=True)
+    tokenized_text = tokenized_text.to(device)
+    output = model(**tokenized_text)
+    return output[0][:, 0, :]
+
+
+class Recommender:
+    SIMILARITY_THRESHOLD = 0.8
+
+    def __init__(self, quotes_df, base_model_name, base_model=AutoModel, ckpt=None):
+        (self.diary_embedder,
+         self.quote_embedder,
+         self.tokenizer) = get_models_and_tokenizer(base_model_name, base_model, ckpt)
+
+        self.quotes = quotes_df['Quote'].to_list()
+        self.authors = quotes_df['Author'].to_list()
+
+        self.quote_embeddings = torch.tensor(np.array(
+            [model_inference(self.quote_embedder, self.tokenizer, q).cpu().detach().numpy() for q in stqdm(self.quotes[:50])]
+        )).squeeze(1)
+
+    def recommend(self, d):
+        d_emb = model_inference(self.diary_embedder, self.tokenizer, d).squeeze().cpu()
+        similarities = F.cosine_similarity(d_emb, self.quote_embeddings, dim=0)
+        above_threshold_indices = (similarities > self.SIMILARITY_THRESHOLD).nonzero().flatten().tolist()
+        if above_threshold_indices:
+            index = random.choice(above_threshold_indices)
+        else:
+            index = torch.argmax(similarities).item()
+        return self.quotes[index], self.authors[index]
+
+
+def get_quote_embeddings(model, tokenizer):
+    quotes = pd.read_csv('quotes-recsys/data/quotes.csv')['Quote'].to_list()
+    return torch.tensor([model_inference(model, tokenizer, q) for q in quotes]).squeeze(1)