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 text_streamer = TextStreamer(tokenizer, skip_special_tokens=True,
12 skip_prompt=True)
13 return model.generate(input_ids=inputs,
14 streamer=text_streamer,
15 max_new_tokens=512)
16
17question = "Does the University of Westminster offer a course on AI, Data and Communication MA?"
18x = query(question, model)