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