Views
No views yet

1<|system|>
2{system_message} <|end|>
3<|user|>
4{Prompt) <|end|>
5<|assistant|>1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2
3model_name_or_path = "thesven/Chatty-McChatterson-3-mini-128k"
4
5# BitsAndBytesConfig for loading the model in 4-bit precision
6bnb_config = BitsAndBytesConfig(
7 load_in_4bit=True,
8 bnb_4bit_quant_type="nf4",
9 bnb_4bit_compute_dtype="float16",
10)
11
12tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
13model = AutoModelForCausalLM.from_pretrained(
14 model_name_or_path,
15 device_map="auto",
16 trust_remote_code=False,
17 revision="main",
18 quantization_config=bnb_config
19)
20model.pad_token = model.config.eos_token_id
21
22prompt_template = '''
23<|user|>
24What is the name of the big tower in Toronto?.<|end|>
25<|assistant|>
26'''
27
28input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
29output = model.generate(inputs=input_ids, temperature=0.1, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=256)
30
31generated_text = tokenizer.decode(output[0, len(input_ids[0]):], skip_special_tokens=True)
32print(generated_text)