27 lines
781 B
Python
27 lines
781 B
Python
"""Pre-training data processor."""
|
|
from typing import Dict, List, Any
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
from pipeline.tokenizer import BpeTokenizer
|
|
from pipeline.processors.base import BaseProcessor
|
|
from pipeline.processors.factory import ProcessorFactory
|
|
|
|
|
|
@ProcessorFactory.register("pt")
|
|
class PreTrainProcessor(BaseProcessor):
|
|
"""Pre-training data processor."""
|
|
|
|
def __init__(self, tokenizer: BpeTokenizer):
|
|
self.tokenizer = tokenizer
|
|
|
|
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
|
segment = input_dict["text"]
|
|
tokens = self.tokenizer.encode(f"{segment}<eos>")
|
|
return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
|
|
|
|
@property
|
|
def output_keys(self) -> List[str]:
|
|
return ["sequence"]
|