Move API header warning to init of pipeline

pipeline.py

@@ -9,17 +9,6 @@ import requests
 import torch
 from transformers import Pipeline
 
-API_HEADERS = {"Accept": "application/vnd.github+json"}
-if os.environ.get("GITHUB_TOKEN") is None:
-    print(
-        "[!] Consider setting GITHUB_TOKEN environment variable 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:
-    API_HEADERS["Authorization"] = f"Bearer {os.environ['GITHUB_TOKEN']}"
-    print("[+] Using GITHUB_TOKEN for authentication")
-
 
 def extract_code_and_docs(text: str):
     """Extract code and documentation from a Python file.
@@ -53,11 +42,11 @@ def extract_code_and_docs(text: str):
     return code_set, docs_set
 
 
-def get_topics(repo_name):
+def get_topics(repo_name, headers=None):
     api_url = f"https://api.github.com/repos/{repo_name}"
     print(f"[+] Getting topics for {repo_name}")
     try:
-        response = requests.get(api_url, headers=API_HEADERS)
+        response = requests.get(api_url, headers=headers)
         response.raise_for_status()
     except requests.exceptions.HTTPError as e:
         print(f"[-] Failed to get topics for {repo_name}: {e}")
@@ -71,19 +60,19 @@ def get_topics(repo_name):
     return topics
 
 
-def download_and_extract(repos):
+def download_and_extract(repos, headers=None):
     extracted_info = {}
     for repo_name in repos:
         extracted_info[repo_name] = {
             "funcs": set(),
             "docs": set(),
-            "topics": get_topics(repo_name),
+            "topics": get_topics(repo_name, headers=headers),
         }
 
         download_url = f"https://api.github.com/repos/{repo_name}/tarball"
         print(f"[+] Extracting functions and docstrings from {repo_name}")
         try:
-            response = requests.get(download_url, headers=API_HEADERS, stream=True)
+            response = requests.get(download_url, headers=headers, stream=True)
             response.raise_for_status()
         except requests.exceptions.HTTPError as e:
             print(f"[-] Failed to download {repo_name}: {e}")
@@ -107,6 +96,20 @@ def download_and_extract(repos):
 
 
 class RepoEmbeddingPipeline(Pipeline):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.API_HEADERS = {"Accept": "application/vnd.github+json"}
+        if os.environ.get("GITHUB_TOKEN") is None:
+            print(
+                "[!] Consider setting GITHUB_TOKEN environment variable 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']}"
+            print("[+] Using GITHUB_TOKEN for authentication")
+
+
     def _sanitize_parameters(self, **kwargs):
         _forward_kwargs = {}
         if "max_length" in kwargs:
@@ -118,7 +121,7 @@ class RepoEmbeddingPipeline(Pipeline):
         if isinstance(inputs, str):
             inputs = (inputs,)
 
-        extracted_infos = download_and_extract(inputs)
+        extracted_infos = download_and_extract(inputs, headers=self.API_HEADERS)
 
         return extracted_infos