Lazyhope commited on
Commit
16c089c
1 Parent(s): 8edac7f

Allow user to manually set github token in each run

Browse files
Files changed (1) hide show
  1. 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.API_HEADERS = {"Accept": "application/vnd.github+json"}
110
- if not github_token:
 
 
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 {}, _forward_kwargs, {}
131
 
132
- def preprocess(self, inputs):
133
  if isinstance(inputs, str):
134
  inputs = [inputs]
135
 
136
- extracted_infos = download_and_extract(inputs, headers=self.API_HEADERS)
 
 
 
 
 
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