feat(#141): embed

Author: h1alexbel

sr-data/src/sr_data/steps/embed.py

@@ -56,7 +56,7 @@ def main(repos, prefix, hf, cohere):
             embed_cohere(cohere, frame, prefix)
         else:
             logger.info(f"Inference checkpoint: {checkpoint}")
-            embeddings = pd.DataFrame(infer(frame["top"].tolist(), checkpoint, hf))
+            embeddings = pd.DataFrame(infer(frame["mcw"].tolist(), checkpoint, hf))
             embeddings.insert(0, 'repo', frame["repo"])
             embeddings.to_csv(f"{prefix}-{model}.csv", index=False)
         logger.info(

sr-data/src/sr_data/steps/extract.py

@@ -61,7 +61,7 @@ def main(repos, out):
     logger.info(
         f"Removed {headingless - len(frame)} repositories that have 0 headings after regex filtering ('{rword}')"
     )
-    frame["top"] = frame["headings"].apply(
+    frame["mcw"] = frame["headings"].apply(
         lambda headings: top_words(headings, 5)
     )
     frame.to_csv(out, index=False)

sr-data/src/tests/test_embed.py

@@ -98,7 +98,7 @@ def test_creates_csv_with_embeddings(self):
         with TemporaryDirectory() as temp:
             main(
                 os.path.join(
-                    os.path.dirname(os.path.realpath(__file__)), "embed.csv"
+                    os.path.dirname(os.path.realpath(__file__)), "to-embed.csv"
                 ),
                 temp,
                 os.environ["HF_TESTING_TOKEN"],

sr-data/src/tests/embed.csv sr-data/src/tests/to-embed.csvsr-data/src/tests/embed.csv renamed to sr-data/src/tests/to-embed.csv

@@ -1,2 +1,2 @@
-repo,top
+repo,mcw
 foo,"['test', 'dummy']"