Show progress bar when running the model

pipeline.py

@@ -6,6 +6,7 @@ from io import BytesIO
 import numpy as np
 import requests
 import torch
+from tqdm import tqdm
 from transformers import Pipeline
 
 
@@ -154,26 +155,37 @@ class RepoEmbeddingPipeline(Pipeline):
 
     def _forward(self, extracted_infos, max_length=512):
         repo_dataset = {}
-        for repo_name, repo_info in extracted_infos.items():
-            entry = {"topics": repo_info.get("topics")}
-
-            print(f"[+] Generating embeddings for {repo_name}")
-            if entry.get("code_embeddings") is None:
-                code_embeddings = [
-                    [func, self.encode(func, max_length).squeeze().tolist()]
-                    for func in repo_info["funcs"]
-                ]
+        num_texts = sum(
+            len(x["funcs"]) + len(x["docs"]) for x in extracted_infos.values()
+        )
+        with tqdm(total=num_texts) as pbar:
+            for repo_name, repo_info in extracted_infos.items():
+                pbar.set_description(f"Processing {repo_name}")
+                entry = {"topics": repo_info.get("topics")}
+
+                print(f"[+] Generating embeddings for {repo_name}")
+
+                code_embeddings = []
+                for func in repo_info["funcs"]:
+                    code_embeddings.append(
+                        [func, self.encode(func, max_length).squeeze().tolist()]
+                    )
+                    pbar.update(1)
+
                 entry["code_embeddings"] = code_embeddings
                 entry["mean_code_embedding"] = (
                     np.mean([x[1] for x in code_embeddings], axis=0).tolist()
                     if code_embeddings
                     else None
                 )
-            if entry.get("doc_embeddings") is None:
-                doc_embeddings = [
-                    [doc, self.encode(doc, max_length).squeeze().tolist()]
-                    for doc in repo_info["docs"]
-                ]
+
+                doc_embeddings = []
+                for doc in repo_info["docs"]:
+                    doc_embeddings.append(
+                        [doc, self.encode(doc, max_length).squeeze().tolist()]
+                    )
+                    pbar.update(1)
+
                 entry["doc_embeddings"] = doc_embeddings
                 entry["mean_doc_embedding"] = (
                     np.mean([x[1] for x in doc_embeddings], axis=0).tolist()
@@ -181,7 +193,7 @@ class RepoEmbeddingPipeline(Pipeline):
                     else None
                 )
 
-            repo_dataset[repo_name] = entry
+                repo_dataset[repo_name] = entry
 
         return repo_dataset