Add application file

app.py

@@ -0,0 +1,109 @@
+import os
+import pickle
+from io import BytesIO
+
+import pandas as pd
+import requests
+import streamlit as st
+
+from inference import retrieve, rerank
+
+
+def get_data(results: pd.DataFrame, data: pd.DataFrame, reranked=False):
+    """Given the corpus indices of the top-k series get the required data for the UI"""
+    if reranked:
+        scores_list = results["cross-score"].tolist()
+    else:
+        scores_list = results.score.tolist()
+
+    titles, scores, covers, urls = [], [], [], []
+    for idx, score in zip(results.corpus_id.tolist(), scores_list):
+        titles.append(data.iloc[idx].romaji)
+        scores.append(score)
+        covers.append(data.iloc[idx].cover)
+        urls.append(data.iloc[idx].url)
+
+    return titles, scores, covers, urls
+
+
+def add_descriptions_to_results(results: pd.DataFrame):
+    """Add the corresponding description to the retrieval results"""
+    idxs = results["corpus_id"].tolist()
+    descs = data.iloc[idxs].input.tolist()
+    results["desc"] = descs
+    return results
+
+
+# Input UI
+st.title("Manga Semantic Search")
+query = st.text_input(
+    "Enter a description of the manga you are searching for:",
+    value="",
+)
+embeddings_path = st.selectbox("Embeddings Corpus", os.listdir("embeddings"))
+top_k = st.number_input(
+    "Number of results", value=5, min_value=1, max_value=100, step=1
+)
+do_rerank = st.checkbox("Re-Rank", value=True)
+k_retrieve = None
+if do_rerank:
+    k_retrieve = st.number_input(
+        "Number of initialy retrieved series",
+        value=50,
+        min_value=1,
+        max_value=500,
+        step=1,
+    )
+
+
+# Convert UI values into the correct function argument values
+model_name = str(embeddings_path).split(".")[-2]
+embeddings_path = os.path.join("embeddings", str(embeddings_path))
+
+
+# Output UI
+if st.button("Search"):
+    if not k_retrieve:
+        k_retrieve = top_k
+
+    # Check that query is not empty
+    if not query:
+        st.write("Please enter a query")
+    # Check that top_k is not > retrieve_k
+    elif top_k > k_retrieve:
+        st.write(
+            "'Number of results' should be less than or equal to 'Number of number of initialy retrieved series'"
+        )
+    else:
+        # Load embedddings and corresponding data table
+        with open(embeddings_path, "rb") as f:
+            data, corpus_embeddings = pickle.load(f).values()
+
+        # Retrieve most similar series
+        results = retrieve(
+            query,
+            corpus_embeddings=corpus_embeddings,
+            model_name=model_name,
+            top_k=int(k_retrieve),
+        )
+        # Re-rank the retrieved series
+        if do_rerank:
+            results = add_descriptions_to_results(results)
+            results = rerank(query, results, top_k=int(top_k))
+
+        # Display results
+        titles, scores, covers, urls = get_data(results, data, do_rerank)
+        for title, score, cover, url in zip(titles, scores, covers, urls):
+            with st.container():
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.markdown(
+                        f"""
+                            ## [{title}]({url})
+                            Score: {score:.2f}
+                        """
+                    )
+                with col2:
+                    response = requests.get(cover)
+                    img = BytesIO(response.content)
+                    st.image(img, width=200)

inference.py

@@ -0,0 +1,62 @@
+import pickle
+import torch
+
+import pandas as pd
+from sentence_transformers import SentenceTransformer, util, CrossEncoder
+
+
+def retrieve(
+    query: str,
+    corpus_embeddings: torch.Tensor,
+    top_k: int = 5,
+    model_name: str = "all-mpnet-base-v2",
+):
+    """Retrieve the most similar series in a corpus given a query"""
+
+    # Embed query
+    model = SentenceTransformer(model_name)
+    prompt_embedding = model.encode(query, convert_to_tensor=True)
+
+    # Find most similar
+    results = util.semantic_search(prompt_embedding, corpus_embeddings, top_k=top_k)[0]
+    results = pd.DataFrame(results, columns=["corpus_id", "score"])
+
+    return results
+
+
+def rerank(
+    query: str,
+    retrieved: pd.DataFrame,
+    top_k: int = 5,
+    model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
+):
+    """Re-rank the retrieved series"""
+
+    # Create pairs of query and descriptions
+    inp = [[query, desc] for desc in retrieved["desc"]]
+
+    # Get scores for each pair
+    cross_encoder = CrossEncoder(model_name)
+    cross_scores = cross_encoder.predict(inp)
+    retrieved["cross-score"] = cross_scores
+
+    # Keep top-k after re-ranking
+    results = retrieved.sort_values("cross-score", ascending=False).iloc[:top_k]
+
+    return results
+
+
+if __name__ == "__main__":
+    with open("embeddings/desc-embeddings.all-mpnet-base-v2.pkl", "rb") as f:
+        data, corpus_embeddings = pickle.load(f).values()
+
+    q = "a series about people battling each other in cooking competitions"
+    results = retrieve(q, corpus_embeddings, top_k=50)
+
+    idxs = results["corpus_id"].tolist()
+    descs = data.iloc[idxs].input.tolist()
+    results["desc"] = descs
+    print(results)
+
+    reranked = rerank(q, results, top_k=5)
+    print(reranked)

requirements.txt

@@ -0,0 +1,4 @@
+pandas==2.0.1
+sentence_transformers==2.2.2
+streamlit==1.22.0
+torch==2.0.0