additional cols and optional device

app_implementation.py

@@ -1,50 +1,61 @@
-import os
 from typing import Dict, List
 
+import torch
 import pandas as pd
-import datasets
 import streamlit as st
-import config
 from findkit import retrieval_pipeline
+
+import config
 from search_utils import (
+    RetrievalPipelineWrapper,
+    get_doc_cols,
     get_repos_with_descriptions,
-    search_f,
-    merge_text_list_cols,
-    setup_retrieval_pipeline,
+    get_retrieval_df,
+    merge_cols,
 )
 
 
 class RetrievalApp:
+    def get_device_options(self):
+        if torch.cuda.is_available:
+            return ["cuda", "cpu"]
+        else:
+            return ["cpu"]
+
+    @st.cache(allow_output_mutation=True)
+    def get_retrieval_df(self):
+        return get_retrieval_df(self.data_path, config.text_list_cols)
+
     def __init__(self, data_path="lambdaofgod/pwc_repositories_with_dependencies"):
+        self.data_path = data_path
+        self.device = st.sidebar.selectbox("device", self.get_device_options())
         print("loading data")
 
-        raw_retrieval_df = (
-            datasets.load_dataset(data_path)["train"]
-            .to_pandas()
-            .drop_duplicates(subset=["repo"])
-            .reset_index(drop=True)
-        )
-        self.retrieval_df = merge_text_list_cols(
-            raw_retrieval_df, config.text_list_cols
-        )
+        self.retrieval_df = self.get_retrieval_df().copy()
 
         model_name = st.sidebar.selectbox("model", config.model_names)
         self.query_encoder_name = "lambdaofgod/query-" + model_name
         self.document_encoder_name = "lambdaofgod/document-" + model_name
 
+        doc_cols = get_doc_cols(model_name)
+
         st.sidebar.text("using models")
         st.sidebar.text("https://huggingface.co/" + self.query_encoder_name)
-        st.sidebar.text("https://huggingface.co/" + self.document_encoder_name)
+        st.sidebar.text("HTTP://huggingface.co/" + self.document_encoder_name)
+
+        self.additional_shown_cols = st.sidebar.multiselect(
+            label="used text features", options=config.text_cols, default=doc_cols
+        )
 
     @staticmethod
     def show_retrieval_results(
-        retrieval_pipe: retrieval_pipeline.RetrievalPipeline,
+        retrieval_pipe: RetrievalPipelineWrapper,
         query: str,
         k: int,
         all_queries: List[str],
         description_length: int,
         repos_by_query: Dict[str, pd.DataFrame],
-        doc_col: str,
+        additional_shown_cols: List[str],
     ):
         print("started retrieval")
         if query in all_queries:
@@ -56,15 +67,14 @@ class RetrievalApp:
                 st.table(get_repos_with_descriptions(retrieval_pipe.X_df, task_repos))
         with st.spinner(text="fetching results"):
             st.write(
-                search_f(retrieval_pipe, query, k, description_length, doc_col).to_html(
+                retrieval_pipe.search(query, k, description_length, additional_shown_cols).to_html(
                     escape=False, index=False
                 ),
                 unsafe_allow_html=True,
             )
         print("finished retrieval")
 
-    @staticmethod
-    def app(retrieval_pipeline, retrieval_df, doc_col):
+    def run_app(self, retrieval_pipeline):
 
         retrieved_results = st.sidebar.number_input("number of results", value=10)
         description_length = st.sidebar.number_input(
@@ -72,17 +82,12 @@ class RetrievalApp:
         )
 
         tasks_deduped = (
-            retrieval_df["tasks"].explode().value_counts().reset_index()
+            self.retrieval_df["tasks"].explode().value_counts().reset_index()
         )  # drop_duplicates().sort_values().reset_index(drop=True)
         tasks_deduped.columns = ["task", "documents per task"]
         with st.sidebar.expander("View test set queries"):
             st.table(tasks_deduped.explode("task"))
-
-        additional_shown_cols = st.sidebar.multiselect(
-            label="additional cols", options=config.text_cols, default=doc_col
-        )
-
-        repos_by_query = retrieval_df.explode("tasks").groupby("tasks")
+        repos_by_query = self.retrieval_df.explode("tasks").groupby("tasks")
         query = st.text_input("input query", value="metric learning")
         RetrievalApp.show_retrieval_results(
             retrieval_pipeline,
@@ -91,16 +96,23 @@ class RetrievalApp:
             tasks_deduped["task"].to_list(),
             description_length,
             repos_by_query,
-            additional_shown_cols,
+            self.additional_shown_cols,
         )
 
-    def main(self):
-        print("setting up retrieval_pipe")
-        doc_col = "dependencies"
-        retrieval_pipeline = setup_retrieval_pipeline(
+    @st.cache(allow_output_mutation=True)
+    def get_retrieval_pipeline(self, displayed_retrieval_df):
+        return RetrievalPipelineWrapper.setup_from_encoder_names(
             self.query_encoder_name,
             self.document_encoder_name,
-            self.retrieval_df[doc_col],
-            self.retrieval_df,
+            displayed_retrieval_df["document"],
+            displayed_retrieval_df,
+            device=self.device,
+        )
+
+    def main(self):
+        print("setting up retrieval_pipe")
+        displayed_retrieval_df = merge_cols(
+            self.retrieval_df.copy(), self.additional_shown_cols
         )
-        RetrievalApp.app(retrieval_pipeline, self.retrieval_df, doc_col)
+        retrieval_pipeline = self.get_retrieval_pipeline(displayed_retrieval_df)
+        self.run_app(retrieval_pipeline)

config.py

@@ -1,11 +1,11 @@
 model_names = [
-    'dependencies-nbow-nbow-mnrl',
-    'readme-nbow-nbow-mnrl',
-    'titles-nbow-nbow-mnrl',
-    'titles#dependencies-nbow-nbow-mnrl',
-    'readme#dependencies-nbow-nbow-mnrl'
+    "dependencies-nbow-nbow-mnrl",
+    "readme-nbow-nbow-mnrl",
+    "titles-nbow-nbow-mnrl",
+    "titles_dependencies-nbow-nbow-mnrl",
+    "readme_dependencies-nbow-nbow-mnrl",
 ]
-best_tasks_path="assets/best_tasks.csv"
-worst_tasks_path="assets/worst_tasks.csv"
+best_tasks_path = "assets/best_tasks.csv"
+worst_tasks_path = "assets/worst_tasks.csv"
 text_cols = ["dependencies", "readme", "titles"]
 text_list_cols = ["titles"]

search_utils.py

@@ -1,7 +1,8 @@
 import os
 from typing import Dict, List
+from dataclasses import dataclass
 
-
+import datasets
 import ast
 import pandas as pd
 import sentence_transformers
@@ -11,6 +12,33 @@ from toolz import partial
 import config
 
 
+def get_doc_cols(model_name):
+    model_name = model_name.replace("query-", "")
+    model_name = model_name.replace("document-", "")
+    return model_name.split("-")[0].split("_")
+
+
+def merge_cols(df, cols):
+    df["document"] = df[cols[0]]
+    for col in cols:
+        df["document"] = df["document"] + " " + df[col]
+    return df
+
+
+def get_retrieval_df(
+    data_path="lambdaofgod/pwc_repositories_with_dependencies", text_list_cols=None
+):
+    raw_retrieval_df = (
+        datasets.load_dataset(data_path)["train"]
+        .to_pandas()
+        .drop_duplicates(subset=["repo"])
+        .reset_index(drop=True)
+    )
+    if text_list_cols:
+        return merge_text_list_cols(raw_retrieval_df, text_list_cols)
+    return raw_retrieval_df
+
+
 def truncate_description(description, length=50):
     return " ".join(description.split()[:length])
 
@@ -19,25 +47,6 @@ def get_repos_with_descriptions(repos_df, repos):
     return repos_df.loc[repos]
 
 
-def search_f(
-    retrieval_pipe: retrieval_pipeline.RetrievalPipeline,
-    query: str,
-    k: int,
-    description_length: int,
-    doc_col: List[str],
-):
-    results = retrieval_pipe.find_similar(query, k)
-    # results['repo'] = results.index
-    results["link"] = "https://github.com/" + results["repo"]
-    for col in doc_col:
-        results[col] = results[col].apply(
-            lambda desc: truncate_description(desc, description_length)
-        )
-    shown_cols = ["repo", "tasks", "link", "distance"]
-    shown_cols = shown_cols + doc_col
-    return results.reset_index(drop=True)[shown_cols]
-
-
 def merge_text_list_cols(retrieval_df, text_list_cols):
     retrieval_df = retrieval_df.copy()
     for col in text_list_cols:
@@ -47,29 +56,50 @@ def merge_text_list_cols(retrieval_df, text_list_cols):
     return retrieval_df
 
 
-def setup_pipeline(
-    extractor: feature_extractors.SentenceEncoderFeatureExtractor,
-    documents_df: pd.DataFrame,
-    text_col: str,
-):
-    retrieval_pipeline.RetrievalPipelineFactory.build(
-        documents_df[text_col], metadata=documents_df
-    )
+@dataclass
+class RetrievalPipelineWrapper:
 
+    pipeline: retrieval_pipeline.RetrievalPipeline
 
-@st.cache(allow_output_mutation=True)
-def setup_retrieval_pipeline(
-    query_encoder_path, document_encoder_path, documents, metadata
-):
-    document_encoder = feature_extractors.SentenceEncoderFeatureExtractor(
-        sentence_transformers.SentenceTransformer(document_encoder_path, device="cpu")
-    )
-    query_encoder = feature_extractors.SentenceEncoderFeatureExtractor(
-        sentence_transformers.SentenceTransformer(query_encoder_path, device="cpu")
-    )
-    retrieval_pipe = retrieval_pipeline.RetrievalPipelineFactory(
-        feature_extractor=document_encoder,
-        query_feature_extractor=query_encoder,
-        index_factory=partial(indexes.NMSLIBIndex.build, distance="cosinesimil"),
-    )
-    return retrieval_pipe.build(documents, metadata=metadata)
+    @classmethod
+    def build_from_encoders(cls, query_encoder, document_encoder, documents, metadata):
+        retrieval_pipe = retrieval_pipeline.RetrievalPipelineFactory(
+            feature_extractor=document_encoder,
+            query_feature_extractor=query_encoder,
+            index_factory=partial(indexes.NMSLIBIndex.build, distance="cosinesimil"),
+        )
+        pipeline = retrieval_pipe.build(documents, metadata=metadata)
+        return RetrievalPipelineWrapper(pipeline)
+
+    def search(
+        self,
+        query: str,
+        k: int,
+        description_length: int,
+        additional_shown_cols: List[str],
+    ):
+        results = self.pipeline.find_similar(query, k)
+        # results['repo'] = results.index
+        results["link"] = "https://github.com/" + results["repo"]
+        for col in additional_shown_cols:
+            results[col] = results[col].apply(
+                lambda desc: truncate_description(desc, description_length)
+            )
+        shown_cols = ["repo", "tasks", "link", "distance"]
+        shown_cols = shown_cols + additional_shown_cols 
+        return results.reset_index(drop=True)[shown_cols]
+
+    @classmethod
+    def setup_from_encoder_names(cls, query_encoder_path, document_encoder_path, documents, metadata, device
+    ):
+        document_encoder = feature_extractors.SentenceEncoderFeatureExtractor(
+            sentence_transformers.SentenceTransformer(
+                document_encoder_path, device=device
+            )
+        )
+        query_encoder = feature_extractors.SentenceEncoderFeatureExtractor(
+            sentence_transformers.SentenceTransformer(query_encoder_path, device=device)
+        )
+        return cls.build_from_encoders(
+            query_encoder, document_encoder, documents, metadata
+        )