from functools import partial import numpy as np import pandas as pd from datasets import load_dataset from tqdm import tqdm from perplexity_lenses.perplexity import KenlmModel def hub_dataset_to_dataframe( path: str, name: str, split: str, sample: int, text_column: str, model: KenlmModel, seed: int = 0, doc_type: str = "Whole document" ) -> pd.DataFrame: load_dataset_fn = partial(load_dataset, path=path) if name: load_dataset_fn = partial(load_dataset_fn, name=name) if split: load_dataset_fn = partial(load_dataset_fn, split=split) dataset = load_dataset_fn(streaming=True).shuffle(buffer_size=10000, seed=seed) if doc_type == "Sentence": dataset = dataset.map(lambda x: [{text_column: sentence, "perplexity": model.get_perplexity(sentence)} for sentence in x[text_column].split("\n")]) else: dataset = dataset.map(lambda x: {text_column: x[text_column], "perplexity": model.get_perplexity(x[text_column])}) instances = [] count = 0 for instance in tqdm(dataset, total=sample): if isinstance(instance, list): for sentence in instance: instances.append(sentence) count += 1 if count == sample: break else: instances.append(instance) count += 1 if count == sample: break return pd.DataFrame(instances) def documents_df_to_sentences_df(df: pd.DataFrame, text_column: str, sample: int, seed: int = 0): df_sentences = pd.DataFrame({text_column: np.array(df[text_column].map(lambda x: x.split("\n")).values.tolist()).flatten()}) return df_sentences.sample(min(sample, df.shape[0]), random_state=seed)