| 15 | |
| 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 = { |
| 37 | "model": "", |
| 38 | "prompt": prompt, |
| 39 | "temperature": 0.5, |
| 40 | "max_new_tokens": 512, |
| 41 | "stop": "</s>", |
| 42 | "stop_token_ids": None, |
| 43 | "echo": False |
| 44 | } |
| 45 | generate_stream_func = generate_stream |
| 46 | output_stream = generate_stream_func(self.model, self.tokenizer, gen_params, "cuda", self.max_sequence_length, force_generate=True) |
| 47 | outputs = self.chatio.return_output(output_stream) |
| 48 | prediction = outputs.strip() |
| 49 | return prediction |
| 50 | |
| 51 | def add_message(self, message): |
| 52 | self.conversation_history.append(message) |
| 53 | |
| 54 | def change_messages(self,messages): |
| 55 | self.conversation_history = messages |
| 56 | |
| 57 | def display_conversation(self, detailed=False): |
| 58 | role_to_color = { |
| 59 | "system": "red", |
| 60 | "user": "green", |
| 61 | "assistant": "blue", |
| 62 | "function": "magenta", |
| 63 | } |
| 64 | print("before_print"+"*"*50) |
| 65 | for message in self.conversation_history: |
| 66 | print_obj = f"{message['role']}: {message['content']} " |
| 67 | if "function_call" in message.keys(): |
| 68 | print_obj = print_obj + f"function_call: {message['function_call']}" |
| 69 | print_obj += "" |
| 70 | print( |
| 71 | colored( |
| 72 | print_obj, |
| 73 | role_to_color[message["role"]], |
| 74 | ) |