Access GitHub token by passing it when initialising a pipeline

pipeline.py

@@ -1,5 +1,4 @@
 import ast
-import os
 import tarfile
 from ast import AsyncFunctionDef, ClassDef, FunctionDef, Module
 from io import BytesIO
@@ -96,20 +95,19 @@ def download_and_extract(repos, headers=None):
 
 
 class RepoEmbeddingPipeline(Pipeline):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, github_token=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.API_HEADERS = {"Accept": "application/vnd.github+json"}
-        if os.environ.get("GITHUB_TOKEN") is None:
+        if not github_token:
             print(
-                "[!] Consider setting GITHUB_TOKEN environment variable to avoid hitting rate limits\n"
+                "[!] 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"
             )
         else:
-            self.API_HEADERS["Authorization"] = f"Bearer {os.environ['GITHUB_TOKEN']}"
+            self.API_HEADERS["Authorization"] = f"Bearer {github_token}"
             print("[+] Using GITHUB_TOKEN for authentication")
 
-
     def _sanitize_parameters(self, **kwargs):
         _forward_kwargs = {}
         if "max_length" in kwargs: