|
from typing import Dict, List, Any
|
|
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
from transformers.generation.logits_process import LogitsProcessorList, InfNanRemoveLogitsProcessor
|
|
from transformers_gad.grammar_utils import IncrementalGrammarConstraint
|
|
from transformers_gad.generation.logits_process import GrammarAlignedOracleLogitsProcessor
|
|
|
|
def safe_int_cast(str, default):
|
|
try:
|
|
return int(str)
|
|
except ValueError:
|
|
return default
|
|
|
|
class EndpointHandler():
|
|
def __init__(self, path=""):
|
|
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
DTYPE = torch.float32
|
|
|
|
self.device = torch.device(DEVICE)
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(path)
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
|
|
self.model = AutoModelForCausalLM.from_pretrained(path)
|
|
self.model.to(self.device)
|
|
self.model.to(dtype=DTYPE)
|
|
self.model.resize_token_embeddings(len(self.tokenizer))
|
|
self.model = torch.compile(self.model, mode='reduce-overhead', fullgraph=True)
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|
|
|
MAX_NEW_TOKENS=512
|
|
MAX_TIME=30
|
|
TEMPERATURE = 1.0
|
|
REPETITION_PENALTY = 1.0
|
|
TOP_P = 1.0
|
|
TOP_K = 0
|
|
|
|
inputs = data.get("inputs", data)
|
|
grammar_str = data.get("grammar", "")
|
|
max_new_tokens = safe_int_cast(data.get("max-new-tokens"), MAX_NEW_TOKENS)
|
|
max_time = safe_int_cast(data.get("max-time"), MAX_TIME)
|
|
|
|
if grammar_str is None or len(grammar_str) == 0 or grammar_str.isspace():
|
|
logits_processors = None
|
|
gad_oracle_processor = None
|
|
else:
|
|
print("=== GOT GRAMMAR ===")
|
|
print(grammar_str)
|
|
print("===================")
|
|
grammar = IncrementalGrammarConstraint(grammar_str, "root", self.tokenizer)
|
|
|
|
|
|
gad_oracle_processor = GrammarAlignedOracleLogitsProcessor(grammar)
|
|
inf_nan_remove_processor = InfNanRemoveLogitsProcessor()
|
|
logits_processors = LogitsProcessorList([
|
|
inf_nan_remove_processor,
|
|
gad_oracle_processor,
|
|
])
|
|
|
|
|
|
input_ids = self.tokenizer.apply_chat_template(
|
|
[{"role": "user", "content": inputs}],
|
|
tokenize=True,
|
|
add_generation_prompt=True,
|
|
return_tensors="pt"
|
|
)
|
|
input_ids = input_ids.to(self.model.device)
|
|
|
|
output = self.model.generate(
|
|
input_ids,
|
|
do_sample=True,
|
|
pad_token_id=self.tokenizer.eos_token_id,
|
|
eos_token_id=self.tokenizer.eos_token_id,
|
|
max_time=max_time,
|
|
max_new_tokens=max_new_tokens,
|
|
top_p=TOP_P,
|
|
top_k=TOP_K,
|
|
repetition_penalty=REPETITION_PENALTY,
|
|
temperature=TEMPERATURE,
|
|
logits_processor=logits_processors,
|
|
num_return_sequences=1,
|
|
return_dict_in_generate=True,
|
|
output_scores=True
|
|
)
|
|
|
|
if gad_oracle_processor is not None:
|
|
gad_oracle_processor.reset()
|
|
|
|
|
|
input_length = 1 if self.model.config.is_encoder_decoder else input_ids.shape[1]
|
|
if (hasattr(output, "sequences")):
|
|
generated_tokens = output.sequences[:, input_length:]
|
|
else:
|
|
generated_tokens = output[:, input_length:]
|
|
|
|
generations = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
|
|
return generations |
