(self, model_name_or_path: str, template:str="tool-llama-single-round", device: str="cuda", cpu_offloading: bool=False, max_sequence_length: int=2048)
| 16 | |
| 17 | class LlamaModel: |
| 18 | def __init__(self, model_name_or_path: str, template:str="tool-llama-single-round", device: str="cuda", cpu_offloading: bool=False, max_sequence_length: int=2048) -> None: |
| 19 | super().__init__() |
| 20 | self.model_name = model_name_or_path |
| 21 | self.template = template |
| 22 | self.max_sequence_length = max_sequence_length |
| 23 | self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False, model_max_length=self.max_sequence_length) |
| 24 | self.model = AutoModelForCausalLM.from_pretrained( |
| 25 | model_name_or_path, low_cpu_mem_usage=True |
| 26 | ) |
| 27 | if self.tokenizer.pad_token_id == None: |
| 28 | self.tokenizer.add_special_tokens({"bos_token": "<s>", "eos_token": "</s>", "pad_token": "<pad>"}) |
| 29 | self.model.resize_token_embeddings(len(self.tokenizer)) |
| 30 | self.use_gpu = (True if device == "cuda" else False) |
| 31 | if (device == "cuda" and not cpu_offloading) or device == "mps": |
| 32 | self.model.to(device) |
| 33 | self.chatio = SimpleChatIO() |
| 34 | |
| 35 | def prediction(self, prompt: str, stop: Optional[List[str]] = None) -> str: |
| 36 | gen_params = { |
nothing calls this directly
no test coverage detected