Views
No views yet
1from transformers import TextStreamer
2
3def chatml(question, model):
4 messages = [{"role": "user", "content": question},]
5
6 inputs = tokenizer.apply_chat_template(messages,
7 tokenize=True,
8 add_generation_prompt=True,
9 return_tensors="pt",).to("cuda")
10
11 print(tokenizer.decode(inputs[0]))
12 text_streamer = TextStreamer(tokenizer, skip_special_tokens=True,
13 skip_prompt=True)
14 return model.generate(input_ids=inputs,
15 streamer=text_streamer,
16 max_new_tokens=512)
17
18
19#Use the following example to test the model:
20question = "Does the University of Westminster offer a course on AI, Data and Communication MA?"
21x = chatml(question, model)
22