Views
No views yet
1!pip install peft
2!pip install transformers
3!pip install bitsandbytes
4!pip install accelerate1# You need a huggingface token that can access llama2
2from huggingface_hub import notebook_login
3notebook_login()1import torch
2from peft import PeftModel, PeftConfig
3from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
4device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5
6bnb_config = BitsAndBytesConfig(
7 load_in_4bit=True,
8 bnb_4bit_use_double_quant=True,
9 bnb_4bit_quant_type="nf4",
10 bnb_4bit_compute_dtype=torch.bfloat16
11)
12
13peft_model_id = "Danjie/SQLMaster_13b"
14config = PeftConfig.from_pretrained(peft_model_id)
15tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
16model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, device_map='auto', quantization_config=bnb_config)
17model.resize_token_embeddings(len(tokenizer) + 1)
18
19# Load the Lora model
20model = PeftModel.from_pretrained(model, peft_model_id)1def create_sql_query(question: str, context: str) -> str:
2 input = "Question: " + question + "\nContext:" + context + "\nAnswer"
3
4 # Encode and move tensor into cuda if applicable.
5 encoded_input = tokenizer(input, return_tensors='pt')
6 encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
7
8 output = model.generate(**encoded_input, max_new_tokens=256)
9 response = tokenizer.decode(output[0], skip_special_tokens=True)
10 response = response[len(input):]
11 return responsecreate_sql_query("What is the highest age of users with name Danjie", "CREATE TABLE user (age INTEGER, name STRING)")