(msg)
| 177 | |
| 178 | |
| 179 | def run_llm(msg): |
| 180 | # vllm |
| 181 | prompt_token_ids = [tokenizer.apply_chat_template(messages, add_generation_prompt=True) for messages in [msg]] |
| 182 | |
| 183 | outputs = llm.generate(prompt_token_ids=prompt_token_ids, sampling_params=sampling_params) |
| 184 | |
| 185 | completion = [output.outputs[0].text for output in outputs] |
| 186 | |
| 187 | return completion[0] |
| 188 | |
| 189 | |
| 190 | # testing for checking |