Returning a dictionary from InformationRetrievalEvaluator instead of writing a CSV

[kevinhu]

InformationRetrievalEvaluator is a convenient way to compute metrics, but it writes to a CSV and only returns a MAP score if write_csv is turned off. Would it be possible to return a dictionary of all computed metrics when write_csv=False?

[tomaarsen]

Hello!

That would indeed be fairly useful. In the meantime, you could use the following:

from sentence_transformers.evaluation import InformationRetrievalEvaluator

class InformationRetrievalEvaluatorWithDict(InformationRetrievalEvaluator):
    def __call__(self, model, output_path: str = None, epoch: int = -1, steps: int = -1, *args, **kwargs) -> float:
        if epoch != -1:
            out_txt = " after epoch {}:".format(epoch) if steps == -1 else " in epoch {} after {} steps:".format(epoch, steps)
        else:
            out_txt = ":"

        logger.info("Information Retrieval Evaluation on " + self.name + " dataset" + out_txt)

        scores = self.compute_metrices(model, *args, **kwargs)

        # Write results to disc
        if output_path is not None and self.write_csv:
            csv_path = os.path.join(output_path, self.csv_file)
            if not os.path.isfile(csv_path):
                fOut = open(csv_path, mode="w", encoding="utf-8")
                fOut.write(",".join(self.csv_headers))
                fOut.write("\n")

            else:
                fOut = open(csv_path, mode="a", encoding="utf-8")

            output_data = [epoch, steps]
            for name in self.score_function_names:
                for k in self.accuracy_at_k:
                    output_data.append(scores[name]['accuracy@k'][k])

                for k in self.precision_recall_at_k:
                    output_data.append(scores[name]['precision@k'][k])
                    output_data.append(scores[name]['recall@k'][k])

                for k in self.mrr_at_k:
                    output_data.append(scores[name]['mrr@k'][k])

                for k in self.ndcg_at_k:
                    output_data.append(scores[name]['ndcg@k'][k])

                for k in self.map_at_k:
                    output_data.append(scores[name]['map@k'][k])

            fOut.write(",".join(map(str, output_data)))
            fOut.write("\n")
            fOut.close()

        return scores

Then you can use InformationRetrievalEvaluatorWithDict instead of InformationRetrievalEvaluator. You can also remove the if output_path ... branch if you're not using write_csv anyways.

• Tom Aarsen

[kevinhu]

This works, thank you! Hope we can see support in the main library soon as well.

[tomaarsen]

I intend to incorporate this into #2449.