Add streamlit widget support to the pipeline

pipeline.py

@@ -6,7 +6,7 @@ from io import BytesIO
 import numpy as np
 import requests
 import torch
-from tqdm import tqdm
+from tqdm.auto import tqdm
 from transformers import Pipeline
 
 
@@ -96,26 +96,38 @@ def download_and_extract(repos, headers=None):
 
 
 class RepoEmbeddingPipeline(Pipeline):
-    def __init__(self, github_token=None, *args, **kwargs):
+    def __init__(self, github_token=None, st_messager=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
+
+        # Streamlit single element container created by st.empty()
+        self.st_messager = st_messager
+
         self.API_HEADERS = {"Accept": "application/vnd.github+json"}
         if not github_token:
-            print(
-                "[!] Consider setting GitHub token to avoid hitting rate limits\n"
-                "For more info, see:"
+            message = (
+                "[*] Consider setting GitHub token to avoid hitting rate limits. \n"
+                "For more info, see: "
                 "https://docs.github.com/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token"
             )
+            print(message)
+            if self.st_messager:
+                self.st_messager.info(message)
         else:
             self.set_github_token(github_token)
 
     def set_github_token(self, github_token):
         self.API_HEADERS["Authorization"] = f"Bearer {github_token}"
-        print("[+] GitHub token set")
+        message = "[+] GitHub token set"
+        print(message)
+        if self.st_messager:
+            self.st_messager.success(message)
 
     def _sanitize_parameters(self, **kwargs):
         _forward_kwargs = {}
         if "max_length" in kwargs:
             _forward_kwargs["max_length"] = kwargs["max_length"]
+        if "st_progress" in kwargs:
+            _forward_kwargs["st_progress"] = kwargs["st_progress"]
 
         return {}, _forward_kwargs, {}
 
@@ -123,6 +135,8 @@ class RepoEmbeddingPipeline(Pipeline):
         if isinstance(inputs, str):
             inputs = (inputs,)
 
+        if self.st_messager:
+            self.st_messager.info("[*] Downloading and extracting repos...")
         extracted_infos = download_and_extract(inputs, headers=self.API_HEADERS)
 
         return extracted_infos
@@ -153,7 +167,7 @@ class RepoEmbeddingPipeline(Pipeline):
 
         return sentence_embeddings
 
-    def _forward(self, extracted_infos, max_length=512):
+    def _forward(self, extracted_infos, max_length=512, st_progress=None):
         repo_dataset = {}
         num_texts = sum(
             len(x["funcs"]) + len(x["docs"]) for x in extracted_infos.values()
@@ -163,14 +177,20 @@ class RepoEmbeddingPipeline(Pipeline):
                 pbar.set_description(f"Processing {repo_name}")
                 entry = {"topics": repo_info.get("topics")}
 
-                print(f"[+] Generating embeddings for {repo_name}")
+                message = f"[*] Generating embeddings for {repo_name}"
+                tqdm.write(message)
+                if self.st_messager:
+                    self.st_messager.info(message)
 
                 code_embeddings = []
                 for func in repo_info["funcs"]:
                     code_embeddings.append(
                         [func, self.encode(func, max_length).squeeze().tolist()]
                     )
+                    
                     pbar.update(1)
+                    if st_progress:
+                        st_progress.progress(pbar.n / pbar.total)
 
                 entry["code_embeddings"] = code_embeddings
                 entry["mean_code_embedding"] = (
@@ -184,7 +204,10 @@ class RepoEmbeddingPipeline(Pipeline):
                     doc_embeddings.append(
                         [doc, self.encode(doc, max_length).squeeze().tolist()]
                     )
+
                     pbar.update(1)
+                    if st_progress:
+                        st_progress.progress(pbar.n / pbar.total)
 
                 entry["doc_embeddings"] = doc_embeddings
                 entry["mean_doc_embedding"] = (