Please use the code below as an example for how to use this model.
1import torch
2from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
3
4def load_model(model_name):
5 # Load tokenizer and model with QLoRA configuration
6 compute_dtype = getattr(torch, 'float16')
7
8 bnb_config = BitsAndBytesConfig(
9 load_in_4bit=True,
10 bnb_4bit_quant_type='nf4',
11 bnb_4bit_compute_dtype=compute_dtype,
12 bnb_4bit_use_double_quant=False,
13 )
14
15 model = AutoModelForCausalLM.from_pretrained(
16 model_name,
17 device_map={"": 0},
18 quantization_config=bnb_config
19 )
20
21
22 # Load Tokenizer
23 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
24 tokenizer.pad_token = tokenizer.eos_token
25 tokenizer.padding_side = "right"
26
27 return model, tokenizer
28
29model, tokenizer = load_model('vagmi/squeal')
30
31prompt = "<s>[INST] Output SQL for the given table structure \n \
32 CREATE TABLE votes (contestant_number VARCHAR, num_votes int); \
33 CREATE TABLE contestants (contestant_number VARCHAR, contestant_name VARCHAR); \
34 What is the contestant number and name of the contestant who got least votes?[/INST]"
35pipe = pipeline(task="text-generation",
36 model=model,
37 tokenizer=tokenizer,
38 max_length=200,
39 device_map='auto', )
40result = pipe(prompt)
41print(result[0]['generated_text'][len(prompt):-1])
Watch me build this model.
Here is the notebook I used to train this model.