From 5231456a3ef834e3585ceec55d51218c71a531d3 Mon Sep 17 00:00:00 2001 From: Carlos Garcia Jurado Suarez Date: Tue, 12 Dec 2023 11:39:08 -0800 Subject: [PATCH] cr feedback to use itertools to flatten --- src/autora/doc/pipelines/main.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/autora/doc/pipelines/main.py b/src/autora/doc/pipelines/main.py index e797ce1..5afc6bf 100644 --- a/src/autora/doc/pipelines/main.py +++ b/src/autora/doc/pipelines/main.py @@ -1,3 +1,4 @@ +import itertools import logging from timeit import default_timer as timer from typing import List @@ -64,7 +65,7 @@ def eval( mlflow.log_text(predictions[i][j], f"prediction_{i}_{j}.txt") # flatten predictions for counting tokens - predictions_flat = [pred for pred_list in predictions for pred in pred_list] + predictions_flat = list(itertools.chain.from_iterable(predictions)) tokens = pred.tokenize(predictions_flat)["input_ids"] total_tokens = sum([len(token) for token in tokens]) mlflow.log_metric("total_tokens", total_tokens)