Allow user to manually set github token in each run
Browse files- pipeline.py +16 -11
pipeline.py
CHANGED
@@ -106,34 +106,39 @@ class RepoEmbeddingPipeline(Pipeline):
|
|
106 |
def __init__(self, github_token=None, *args, **kwargs):
|
107 |
super().__init__(*args, **kwargs)
|
108 |
|
109 |
-
self.
|
110 |
-
if
|
|
|
|
|
111 |
print(
|
112 |
"[*] Consider setting GitHub token to avoid hitting rate limits. \n"
|
113 |
"For more info, see: "
|
114 |
"https://docs.github.com/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token"
|
115 |
)
|
116 |
-
else:
|
117 |
-
self.set_github_token(github_token)
|
118 |
-
|
119 |
-
def set_github_token(self, github_token):
|
120 |
-
self.API_HEADERS["Authorization"] = f"Bearer {github_token}"
|
121 |
-
print("[+] GitHub token set")
|
122 |
|
123 |
def _sanitize_parameters(self, **kwargs):
|
|
|
|
|
|
|
|
|
124 |
_forward_kwargs = {}
|
125 |
if "max_length" in kwargs:
|
126 |
_forward_kwargs["max_length"] = kwargs["max_length"]
|
127 |
if "st_progress" in kwargs:
|
128 |
_forward_kwargs["st_progress"] = kwargs["st_progress"]
|
129 |
|
130 |
-
return
|
131 |
|
132 |
-
def preprocess(self, inputs):
|
133 |
if isinstance(inputs, str):
|
134 |
inputs = [inputs]
|
135 |
|
136 |
-
|
|
|
|
|
|
|
|
|
|
|
137 |
|
138 |
return extracted_infos
|
139 |
|
|
|
106 |
def __init__(self, github_token=None, *args, **kwargs):
|
107 |
super().__init__(*args, **kwargs)
|
108 |
|
109 |
+
self.github_token = github_token
|
110 |
+
if self.github_token:
|
111 |
+
print("[+] GitHub token set!")
|
112 |
+
else:
|
113 |
print(
|
114 |
"[*] Consider setting GitHub token to avoid hitting rate limits. \n"
|
115 |
"For more info, see: "
|
116 |
"https://docs.github.com/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token"
|
117 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
118 |
|
119 |
def _sanitize_parameters(self, **kwargs):
|
120 |
+
preprocess_kwargs = {}
|
121 |
+
if "github_token" in kwargs:
|
122 |
+
preprocess_kwargs["github_token"] = kwargs["github_token"]
|
123 |
+
|
124 |
_forward_kwargs = {}
|
125 |
if "max_length" in kwargs:
|
126 |
_forward_kwargs["max_length"] = kwargs["max_length"]
|
127 |
if "st_progress" in kwargs:
|
128 |
_forward_kwargs["st_progress"] = kwargs["st_progress"]
|
129 |
|
130 |
+
return preprocess_kwargs, _forward_kwargs, {}
|
131 |
|
132 |
+
def preprocess(self, inputs, github_token=None):
|
133 |
if isinstance(inputs, str):
|
134 |
inputs = [inputs]
|
135 |
|
136 |
+
headers = {"Accept": "application/vnd.github+json"}
|
137 |
+
token = github_token or self.github_token
|
138 |
+
if token:
|
139 |
+
headers["Authorization"] = f"Bearer {token}"
|
140 |
+
|
141 |
+
extracted_infos = download_and_extract(inputs, headers=headers)
|
142 |
|
143 |
return extracted_infos
|
144 |
|