File size: 1,003 Bytes
bb48904 c12f2c8 bb48904 c12f2c8 bb48904 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from expand import *
from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast, BatchEncoding
from dataclasses import dataclass
from completions import prepare_inputs, find_next_tokens
type Tokenizer = PreTrainedTokenizer | PreTrainedTokenizerFast
@dataclass
class ExpanderOneBatchLLM:
model: PreTrainedModel
tokenizer: Tokenizer
def expand(self, batch: Batch) -> ExpansionOneResultBatch:
inputs = prepare_inputs([s.get_all_tokens() for s in batch.items], self.tokenizer, self.model.device)
next_tokens = find_next_tokens(self.model, inputs, self.tokenizer)
results = []
for s, next_tokens in zip(batch.items, next_tokens):
expansions = [Expansion(token=token, cost=cost) for token, cost in next_tokens if cost + s.get_remaining_budget() >= 0]
results.append(ExpansionOneResult(series=s, expansions=expansions))
return ExpansionOneResultBatch(items=results)
|