Views
No views yet

1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2
3model_name_or_path = "thesven/Llama3-8B-SFT-SyntheticMedical-bnb-4bit"
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<|begin_of_text|><|start_header_id|>system<|end_header_id|>
24
25You are an expert in the field of anatomy, help explain its topics to me.<|eot_id|><|start_header_id|>user<|end_header_id|>
26
27What is the function of the hamstring?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
28'''
29
30input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
31output = model.generate(inputs=input_ids, temperature=0.1, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
32
33print(generated_text)
34