app refactor and new models

app_implementation.py

@@ -0,0 +1,106 @@
+import os
+from typing import Dict, List
+
+import pandas as pd
+import datasets
+import streamlit as st
+import config
+from findkit import retrieval_pipeline
+from search_utils import (
+    get_repos_with_descriptions,
+    search_f,
+    merge_text_list_cols,
+    setup_retrieval_pipeline,
+)
+
+
+class RetrievalApp:
+    def __init__(self, data_path="lambdaofgod/pwc_repositories_with_dependencies"):
+        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
+        )
+
+        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
+
+        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)
+
+    @staticmethod
+    def show_retrieval_results(
+        retrieval_pipe: retrieval_pipeline.RetrievalPipeline,
+        query: str,
+        k: int,
+        all_queries: List[str],
+        description_length: int,
+        repos_by_query: Dict[str, pd.DataFrame],
+        doc_col: str,
+    ):
+        print("started retrieval")
+        if query in all_queries:
+            with st.expander(
+                "query is in gold standard set queries. Toggle viewing gold standard results?"
+            ):
+                st.write("gold standard results")
+                task_repos = repos_by_query.get_group(query)
+                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(
+                    escape=False, index=False
+                ),
+                unsafe_allow_html=True,
+            )
+        print("finished retrieval")
+
+    @staticmethod
+    def app(retrieval_pipeline, retrieval_df, doc_col):
+
+        retrieved_results = st.sidebar.number_input("number of results", value=10)
+        description_length = st.sidebar.number_input(
+            "number of used description words", value=10
+        )
+
+        tasks_deduped = (
+            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")
+        query = st.text_input("input query", value="metric learning")
+        RetrievalApp.show_retrieval_results(
+            retrieval_pipeline,
+            query,
+            retrieved_results,
+            tasks_deduped["task"].to_list(),
+            description_length,
+            repos_by_query,
+            additional_shown_cols,
+        )
+
+    def main(self):
+        print("setting up retrieval_pipe")
+        doc_col = "dependencies"
+        retrieval_pipeline = setup_retrieval_pipeline(
+            self.query_encoder_name,
+            self.document_encoder_name,
+            self.retrieval_df[doc_col],
+            self.retrieval_df,
+        )
+        RetrievalApp.app(retrieval_pipeline, self.retrieval_df, doc_col)

config.py

@@ -1,4 +1,11 @@
-query_encoder_model_name = "lambdaofgod/query_nbow_embedder"
-document_encoder_model_name = "lambdaofgod/document_nbow_embedder"
+model_names = [
+    '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"
+text_cols = ["dependencies", "readme", "titles"]
+text_list_cols = ["titles"]

pages/1_Retrieval_App.py

@@ -1,151 +1,5 @@
-import os
-from typing import Dict, List
+from app_implementation import RetrievalApp
 
-import datasets
-import pandas as pd
-import sentence_transformers
-import streamlit as st
-from findkit import feature_extractors, indexes, retrieval_pipeline
-from toolz import partial
-import config
 
-
-def truncate_description(description, length=50):
-    return " ".join(description.split()[:length])
-
-
-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 show_retrieval_results(
-    retrieval_pipe: retrieval_pipeline.RetrievalPipeline,
-    query: str,
-    k: int,
-    all_queries: List[str],
-    description_length: int,
-    repos_by_query: Dict[str, pd.DataFrame],
-    doc_col: str,
-):
-    print("started retrieval")
-    if query in all_queries:
-        with st.expander(
-            "query is in gold standard set queries. Toggle viewing gold standard results?"
-        ):
-            st.write("gold standard results")
-            task_repos = repos_by_query.get_group(query)
-            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(
-                escape=False, index=False
-            ),
-            unsafe_allow_html=True,
-        )
-    print("finished retrieval")
-
-
-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
-    )
-
-
-@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)
-
-
-def app(retrieval_pipeline, retrieval_df, doc_col):
-
-    retrieved_results = st.sidebar.number_input("number of results", value=10)
-    description_length = st.sidebar.number_input(
-        "number of used description words", value=10
-    )
-
-    tasks_deduped = (
-        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=[doc_col], default=doc_col
-    )
-
-    repos_by_query = retrieval_df.explode("tasks").groupby("tasks")
-    query = st.text_input("input query", value="metric learning")
-    show_retrieval_results(
-        retrieval_pipeline,
-        query,
-        retrieved_results,
-        tasks_deduped["task"].to_list(),
-        description_length,
-        repos_by_query,
-        additional_shown_cols,
-    )
-
-
-def app_main(
-    query_encoder_path,
-    document_encoder_path,
-    data_path,
-):
-    print("loading data")
-
-    retrieval_df = (
-        datasets.load_dataset(data_path)["train"]
-        .to_pandas()
-        .drop_duplicates(subset=["repo"])
-        .reset_index(drop=True)
-    )
-    print("setting up retrieval_pipe")
-    doc_col = "dependencies"
-    retrieval_pipeline = setup_retrieval_pipeline(
-        query_encoder_path, document_encoder_path, retrieval_df[doc_col], retrieval_df
-    )
-    app(retrieval_pipeline, retrieval_df, doc_col)
-
-
-app_main(
-    query_encoder_path=config.query_encoder_model_name,
-    document_encoder_path=config.document_encoder_model_name,
-    data_path="lambdaofgod/pwc_repositories_with_dependencies",
-)
+app = RetrievalApp()
+app.main()

search_utils.py

@@ -0,0 +1,75 @@
+import os
+from typing import Dict, List
+
+
+import ast
+import pandas as pd
+import sentence_transformers
+import streamlit as st
+from findkit import feature_extractors, indexes, retrieval_pipeline
+from toolz import partial
+import config
+
+
+def truncate_description(description, length=50):
+    return " ".join(description.split()[:length])
+
+
+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:
+        retrieval_df[col] = retrieval_df[col].apply(
+            lambda t: " ".join(ast.literal_eval(t))
+        )
+    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
+    )
+
+
+@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)