Spaces:
Build error
Build error
File size: 4,617 Bytes
d660b02 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
from sklearn.model_selection import train_test_split
from llm_engineering.application.preprocessing.operations.chunking import chunk_document
from llm_engineering.domain.cleaned_documents import CleanedDocument
from llm_engineering.domain.dataset import (
InstructDataset,
InstructDatasetSample,
InstructTrainTestSplit,
PreferenceDataset,
PreferenceDatasetSample,
PreferenceTrainTestSplit,
)
from llm_engineering.domain.types import DataCategory
def create_instruct_train_test_split(
data: dict[DataCategory, InstructDataset], test_size=0.2, random_state=42
) -> InstructTrainTestSplit:
train_data = {}
test_data = {}
for category, dataset in data.items():
samples = dataset.samples
samples_dicts = [sample.model_dump() for sample in samples]
if len(samples_dicts) > 0:
train_samples_dicts, test_samples_dicts = train_test_split(
samples_dicts, test_size=test_size, random_state=random_state
)
train_samples = [InstructDatasetSample(**sample_dict) for sample_dict in train_samples_dicts]
test_samples = [InstructDatasetSample(**sample_dict) for sample_dict in test_samples_dicts]
else:
train_samples = []
test_samples = []
train_dataset = InstructDataset(category=category, samples=train_samples)
test_dataset = InstructDataset(category=category, samples=test_samples)
train_data[category] = train_dataset
test_data[category] = test_dataset
return InstructTrainTestSplit(train=train_data, test=test_data, test_split_size=test_size)
def create_preference_train_test_split(
data: dict[DataCategory, PreferenceDataset], test_size=0.2, random_state=42
) -> PreferenceTrainTestSplit:
train_data = {}
test_data = {}
for category, dataset in data.items():
samples = dataset.samples
samples_dicts = [sample.model_dump() for sample in samples]
if len(samples_dicts) > 0:
train_samples_dicts, test_samples_dicts = train_test_split(
samples_dicts, test_size=test_size, random_state=random_state
)
train_samples = [PreferenceDatasetSample(**sample_dict) for sample_dict in train_samples_dicts]
test_samples = [PreferenceDatasetSample(**sample_dict) for sample_dict in test_samples_dicts]
else:
train_samples = []
test_samples = []
train_dataset = PreferenceDataset(category=category, samples=train_samples)
test_dataset = PreferenceDataset(category=category, samples=test_samples)
train_data[category] = train_dataset
test_data[category] = test_dataset
return PreferenceTrainTestSplit(train=train_data, test=test_data, test_split_size=test_size)
def filter_short_answers(
data: dict[DataCategory, PreferenceDataset], min_length: int = 100
) -> dict[DataCategory, PreferenceDataset]:
def is_long_enough(example: PreferenceDatasetSample) -> bool:
return len(example.chosen) >= min_length
filtered_data = {}
for category, dataset in data.items():
filetered_dataset_samples = list(filter(is_long_enough, dataset.samples))
filtered_dataset = PreferenceDataset(category=category, samples=filetered_dataset_samples)
filtered_data[category] = filtered_dataset
return filtered_data
def filter_answer_format(data: dict[DataCategory, PreferenceDataset]) -> dict[DataCategory, PreferenceDataset]:
def is_valid_format(example: PreferenceDatasetSample) -> bool:
chosen = example.chosen
return len(chosen) > 0 and chosen[0].isupper() and chosen[-1] in (".", "!", "?")
filtered_data = {}
for category, dataset in data.items():
filetered_dataset_samples = list(filter(is_valid_format, dataset.samples))
filtered_dataset = PreferenceDataset(category=category, samples=filetered_dataset_samples)
filtered_data[category] = filtered_dataset
return filtered_data
def extract_substrings(
documents: list[CleanedDocument], min_length: int = 1000, max_length: int = 2000
) -> list[CleanedDocument]:
extracts = []
for document in documents:
document_extracts = chunk_document(document.content, min_length, max_length)
for extract in document_extracts:
subdocument = document.model_copy()
subdocument.content = extract
extracts.append(subdocument)
return extracts
|