# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import csv
import pandas as pd
PARQUET_EXTENSION = ".parquet"
[docs]
def save_infer_prob(output_filepath: str, probabilities: list[list[list[float]]]):
"""Save batched probabilities into a parquet file."""
df_probs = pd.DataFrame(probabilities)
df_probs.to_parquet(f"{output_filepath}{PARQUET_EXTENSION}")
[docs]
def load_infer_prob(input_filepath: str) -> list[list[list[float]]]:
"""Retrieve batched probabilities from a parquet file."""
probs_count_in_first_batch = None
def to_list(probs):
"""Ensure number of probabilities is the same for all entries."""
probs_list = list(probs)
nonlocal probs_count_in_first_batch
probs_count_in_first_batch = probs_count_in_first_batch or len(probs_list)
if probs_count_in_first_batch != len(probs_list):
raise ValueError(
f"Reading `{input_filepath}{PARQUET_EXTENSION}`: inconsistent number of"
f"probs across entries: len({probs_list})!={probs_count_in_first_batch}"
)
return probs_list
df_probs = pd.read_parquet(f"{input_filepath}{PARQUET_EXTENSION}")
probabilities = df_probs.to_numpy().tolist()
probabilities = [[to_list(probs) for probs in batch] for batch in probabilities]
return probabilities
# The inference probabilities (`probabilities`) are structured as follows:
# (the example below assumes 4 batches of batch_size=2 and, for each of these,
# 4 probabilities corresponding to the multiple choices A, B, C, D)
#
# [
# [ <-- batch no 0:
# [p_0_0_A, p_0_0_B, p_0_0_C, p_0_0_D], <-- batch index = 0
# [p_0_1_A, p_0_1_B, p_0_1_C, p_0_1_D], <-- batch index = 1
# ],
# [ <-- batch no 1:
# [p_1_0_A, p_1_0_B, p_1_0_C, p_1_0_D], <-- batch index = 0
# [p_1_1_A, p_1_1_B, p_1_1_C, p_1_1_D], <-- batch index = 1
# ],
# [ <-- batch no 2:
# [p_2_0_A, p_2_0_B, p_2_0_C, p_2_0_D], <-- batch index = 0
# [p_2_1_A, p_2_1_B, p_2_1_C, p_2_1_D], <-- batch index = 1
# ],
# [ <-- batch no 3:
# [p_3_0_A, p_3_0_B, p_3_0_C, p_3_0_D], <-- batch index = 0
# [p_3_1_A, p_3_1_B, p_3_1_C, p_3_1_D], <-- batch index = 1
# ]
# ]
#
# We save these into a .csv file of the following format:
# - Every row corresponds to a batch.
# - Within each row, the batch items are strings separated by comma (,).
# - Each item (string) contains a list of probabilities (floats).
#
# batch index = 0 batch index = 1 batch no
# <--------------------------------> , <--------------------------------> |
# "[p_0_0_A, p_0_0_B, p_0_0_C, p_0_0_D]","[p_0_1_A, p_0_1_B, p_0_1_C, p_0_1_D]" <--0
# "[p_1_0_A, p_1_0_B, p_1_0_C, p_1_0_D]","[p_1_1_A, p_1_1_B, p_1_1_C, p_1_1_D]" <--1
# "[p_2_0_A, p_2_0_B, p_2_0_C, p_2_0_D]","[p_2_1_A, p_2_1_B, p_2_1_C, p_2_1_D]" <--2
# "[p_3_0_A, p_3_0_B, p_3_0_C, p_3_0_D]","[p_3_1_A, p_3_1_B, p_3_1_C, p_3_1_D]" <--3
#
[docs]
def save_infer_prob_csv(output_filepath: str, probabilities: list[list[list[float]]]):
"""Save batched probabilities into a csv file."""
with open(output_filepath, "w") as write_obj:
csv_writer = csv.writer(write_obj)
csv_writer.writerows(probabilities)
[docs]
def load_infer_prob_csv(input_filepath: str) -> list[list[list[float]]]:
"""Retrieve batched probabilities from a csv file."""
probs_count_in_first_batch = None
try:
with open(input_filepath) as read_obj:
csv_reader = csv.reader(read_obj)
probabilities = []
for batch in csv_reader:
probabilities_batch = []
for entry in batch:
probs_list = str_to_float_list(entry)
# Number of probabilities must be the same for all entries.
probs_count_in_first_batch = probs_count_in_first_batch or len(
probs_list
)
if probs_count_in_first_batch != len(probs_list):
raise ValueError(
f"Reading {input_filepath}: inconsistent number of probs "
f"across entries: len({probs_list}) != "
f"{probs_count_in_first_batch}"
)
probabilities_batch.append(probs_list)
probabilities.append(probabilities_batch)
return probabilities
except FileNotFoundError:
raise FileNotFoundError(f"{load_infer_prob}: Path {input_filepath} not found!")
[docs]
def str_to_float_list(input: str) -> list[float]:
"""Convert an `str` representing a list of `floats` to an actual list of `floats`.
Example: input: `[1.1, 2.2, 3.3]` => output: [1.1, 2.2, 3.3]
"""
# 1) Get rid of '[' and ']'.
if (input[0] != "[") or (input[-1] != "]"):
raise ValueError(
f"Input `{input}` must start with '[' and end with ']' to represent a list"
)
input = input[1:-1]
# 2) Convert string to a list of items.
list_of_items = input.split(", ")
if not len(list_of_items):
raise ValueError(f"List `{list_of_items}` does NOT contain any items")
# 3) Cast all list items to `float`.
try:
list_of_floats = [float(item) for item in list_of_items]
except ValueError:
raise ValueError(f"List `{list_of_items}` should contain probabilities")
return list_of_floats
