File size: 1,384 Bytes
4d061f7 |
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 |
from typing import Dict, Any
from transformers import PreTrainedTokenizer
# Define la función de tokenización
def tokenize(element: Dict[str, Any], tokenizer: PreTrainedTokenizer, context_length: int) -> Dict[str, Any]:
"""
Tokeniza un elemento de texto usando el tokenizador especificado.
Args:
element (Dict[str, Any]): Un diccionario con la clave "text" que contiene el texto a tokenizar.
tokenizer (PreTrainedTokenizer): Un tokenizador de Hugging Face usado para tokenizar el texto.
context_length (int): Tamaño máximo de los fragmentos de texto.
Returns:
Dict[str, Any]: Un diccionario que contiene los ids de los tokens en la clave "input_ids".
"""
outputs = tokenizer(
element["text"],
truncation=False,
max_length=context_length,
return_overflowing_tokens=True,
return_length=True,
)
# Combinar todos los tokens de los documentos tokenizados
combined = []
for tokenized_doc in outputs['input_ids']:
combined.extend(tokenized_doc + [tokenizer.eos_token_id]) # Usar extend() para mayor eficiencia
# Dividir el texto tokenizado en segmentos del tamaño de context_length
input_batch = [
combined[i:i + context_length] for i in range(0, len(combined) - context_length, context_length)
]
return {"input_ids": input_batch}
|