Batch inference provider for candidate generation models.
| 29 | |
| 30 | |
| 31 | class CandidateBatchInferenceProvider(object): |
| 32 | """ |
| 33 | Batch inference provider for candidate generation models. |
| 34 | """ |
| 35 | |
| 36 | def __init__(self, model_path, prompt_template_name="alpaca") -> None: |
| 37 | super().__init__() |
| 38 | self.template = prompt_templates[prompt_template_name] |
| 39 | try: |
| 40 | tokenizer = AutoTokenizer.from_pretrained(model_path) |
| 41 | except: |
| 42 | tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False) |
| 43 | model = AutoModelForCausalLM.from_pretrained( |
| 44 | model_path, |
| 45 | torch_dtype=torch.bfloat16, |
| 46 | device_map="auto", |
| 47 | ) |
| 48 | if tokenizer.pad_token is None: |
| 49 | self.smart_tokenizer_and_embedding_resize( |
| 50 | special_tokens_dict=dict(pad_token="[PAD]"), |
| 51 | tokenizer=tokenizer, |
| 52 | model=model, |
| 53 | ) |
| 54 | if "llama" in model_path: |
| 55 | tokenizer.add_special_tokens( |
| 56 | { |
| 57 | "eos_token": "</s>", |
| 58 | "bos_token": "</s>", |
| 59 | "unk_token": "</s>", |
| 60 | } |
| 61 | ) |
| 62 | self.tokenizer = tokenizer |
| 63 | |
| 64 | model.config.pad_token_id = self.tokenizer.pad_token_id = 0 # unk |
| 65 | model.config.bos_token_id = 1 |
| 66 | model.config.eos_token_id = 2 |
| 67 | model.eval() |
| 68 | |
| 69 | if torch.__version__ >= "2" and sys.platform != "win32": |
| 70 | model = torch.compile(model) |
| 71 | |
| 72 | self.model = model |
| 73 | self.prepared = [] |
| 74 | self.pattern = re.compile( |
| 75 | r"<unk>|<pad>|<s>|</s>|\[PAD\]|<\|endoftext\|>|\[UNK\]|\[CLS\]|\[MASK\]|<\|startofpiece\|>|<\|endofpiece\|>|\[gMASK\]|\[sMASK\]" |
| 76 | ) |
| 77 | |
| 78 | def generate_prompt(self, instruction, input=None, label=None): |
| 79 | if input: |
| 80 | res = self.template["prompt_input"].format( |
| 81 | instruction=instruction, input=input |
| 82 | ) |
| 83 | else: |
| 84 | res = self.template["prompt_no_input"].format(instruction=instruction) |
| 85 | if label: |
| 86 | res = f"{res}{label}" |
| 87 | return res |
| 88 |
no outgoing calls
no test coverage detected