| 41 | self.chatio = SimpleChatIO() |
| 42 | |
| 43 | def prediction(self, prompt: str, stop: Optional[List[str]] = None) -> str: |
| 44 | with torch.no_grad(): |
| 45 | gen_params = { |
| 46 | "model": "", |
| 47 | "prompt": prompt, |
| 48 | "temperature": 0.5, |
| 49 | "max_new_tokens": 512, |
| 50 | "stop": "</s>", |
| 51 | "stop_token_ids": None, |
| 52 | "echo": False |
| 53 | } |
| 54 | generate_stream_func = generate_stream |
| 55 | output_stream = generate_stream_func(self.model, self.tokenizer, gen_params, "cuda", self.max_sequence_length, force_generate=True) |
| 56 | outputs = self.chatio.return_output(output_stream) |
| 57 | prediction = outputs.strip() |
| 58 | return prediction |
| 59 | |
| 60 | def add_message(self, message): |
| 61 | self.conversation_history.append(message) |