Access GitHub token by passing it when initialising a pipeline
Browse files- pipeline.py +4 -6
pipeline.py
CHANGED
@@ -1,5 +1,4 @@
|
|
1 |
import ast
|
2 |
-
import os
|
3 |
import tarfile
|
4 |
from ast import AsyncFunctionDef, ClassDef, FunctionDef, Module
|
5 |
from io import BytesIO
|
@@ -96,20 +95,19 @@ def download_and_extract(repos, headers=None):
|
|
96 |
|
97 |
|
98 |
class RepoEmbeddingPipeline(Pipeline):
|
99 |
-
def __init__(self, *args, **kwargs):
|
100 |
super().__init__(*args, **kwargs)
|
101 |
self.API_HEADERS = {"Accept": "application/vnd.github+json"}
|
102 |
-
if
|
103 |
print(
|
104 |
-
"[!] Consider setting GITHUB_TOKEN
|
105 |
"For more info, see:"
|
106 |
"https://docs.github.com/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token"
|
107 |
)
|
108 |
else:
|
109 |
-
self.API_HEADERS["Authorization"] = f"Bearer {
|
110 |
print("[+] Using GITHUB_TOKEN for authentication")
|
111 |
|
112 |
-
|
113 |
def _sanitize_parameters(self, **kwargs):
|
114 |
_forward_kwargs = {}
|
115 |
if "max_length" in kwargs:
|
|
|
1 |
import ast
|
|
|
2 |
import tarfile
|
3 |
from ast import AsyncFunctionDef, ClassDef, FunctionDef, Module
|
4 |
from io import BytesIO
|
|
|
95 |
|
96 |
|
97 |
class RepoEmbeddingPipeline(Pipeline):
|
98 |
+
def __init__(self, github_token=None, *args, **kwargs):
|
99 |
super().__init__(*args, **kwargs)
|
100 |
self.API_HEADERS = {"Accept": "application/vnd.github+json"}
|
101 |
+
if not github_token:
|
102 |
print(
|
103 |
+
"[!] Consider setting GITHUB_TOKEN to avoid hitting rate limits\n"
|
104 |
"For more info, see:"
|
105 |
"https://docs.github.com/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token"
|
106 |
)
|
107 |
else:
|
108 |
+
self.API_HEADERS["Authorization"] = f"Bearer {github_token}"
|
109 |
print("[+] Using GITHUB_TOKEN for authentication")
|
110 |
|
|
|
111 |
def _sanitize_parameters(self, **kwargs):
|
112 |
_forward_kwargs = {}
|
113 |
if "max_length" in kwargs:
|