SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[example/glue] fix compute_metrics_fn for bart like models (#7248)

Browse Source 
* fix compute_metrics_fn

* p.predictions -> preds

* apply suggestions
This commit is contained in:
Suraj Patil
2020-09-21 15:04:20 +05:30
committed by GitHub
parent 39062d05f0
commit 43b9d93875
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
examples/text-classification/run_glue.py
View File

@@ -150,10 +150,11 @@ def main():
def build_compute_metrics_fn(task_name: str) -> Callable[[EvalPrediction], Dict]:
def compute_metrics_fn(p: EvalPrediction):
preds = p.predictions[0] if type(p.predictions) == tuple else p.predictions
if output_mode == "classification":
preds = np.argmax(p.predictions, axis=1)
elif output_mode == "regression":
preds = np.squeeze(p.predictions)
preds = np.argmax(preds, axis=1)
else: # regression
preds = np.squeeze(preds)
return glue_compute_metrics(task_name, preds, p.label_ids)
return compute_metrics_fn