| 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) |