Views
No views yet
1import sys
2import time
3
4import torch
5from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
6
7
8prompt_template = """
9<s>[INST] <<SYS>>
10You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
11
12If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
13<</SYS>>
14
15{question} [/INST]
16{response}
17"""
18
19
20def ask_question(model, tokenizer, question):
21 pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_new_tokens=2048)
22 prompt = prompt_template.format(question=question, response="")
23 tokens_in = len(tokenizer(prompt)["input_ids"])
24 start = time.time()
25 result = pipe(prompt)
26 end = time.time()
27 generated_text = result[0]['generated_text']
28 tokens_out = len(tokenizer(generated_text)["input_ids"])
29 print(generated_text)
30 tokens_generated = tokens_out - tokens_in
31 time_taken = end - start
32 tokens_per_second = tokens_generated / time_taken
33 print(f"{tokens_generated} tokens in {time_taken:.2f}s: {tokens_per_second:.2f} tokens/s)")
34
35
36
37def test_model():
38 model_name = "Qwen1.5-0.5B-openassistant-guanaco-llama2-format"
39 tokenizer = AutoTokenizer.from_pretrained(model_name)
40 model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda", torch_dtype=torch.bfloat16)
41
42 question = input("You: ")
43 ask_question(model, tokenizer, question)
44
45
46
47if __name__ == "__main__":
48 test_model()