Views
No views yet
pip install -q transformers==4.35.0 torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 optimum==1.13.2 auto-gptq==0.4.21from transformers import AutoTokenizer
2from auto_gptq import AutoGPTQForCausalLM
3
4model_name = 'support-pvelocity/Code-Llama-2-13B-instruct-text2sql-GPTQ'
5
6model = AutoGPTQForCausalLM.from_quantized(model_name, use_safetensors=True, device_map='auto')
7tokenizer = AutoTokenizer.from_pretrained(model_name)
8
9table = "CREATE TABLE sales ( sale_id number PRIMARY KEY, product_id number, customer_id number, salesperson_id number, sale_date DATE, quantity number, FOREIGN KEY (product_id) REFERENCES products(product_id), FOREIGN KEY (customer_id) REFERENCES customers(customer_id), FOREIGN KEY (salesperson_id) REFERENCES salespeople(salesperson_id)); CREATE TABLE product_suppliers ( supplier_id number PRIMARY KEY, product_id number, supply_price number, FOREIGN KEY (product_id) REFERENCES products(product_id)); CREATE TABLE customers ( customer_id number PRIMARY KEY, name text, address text ); CREATE TABLE salespeople ( salesperson_id number PRIMARY KEY, name text, region text ); CREATE TABLE product_suppliers ( supplier_id number PRIMARY KEY, product_id number, supply_price number );"
10
11question = 'Find the salesperson who made the most sales.'
12
13prompt = f"[INST] Write SQLite query to answer the following question given the database schema. Please wrap your code answer using ```: Schema: {table} Question: {question} [/INST] Here is the SQLite query to answer to the question: {question}: ``` "
14
15tokens = tokenizer(prompt, return_tensors="pt").to('cuda:0')
16input_ids = tokens.input_ids
17
18generated_ids = model.generate(input_ids=input_ids, max_length=4048, pad_token_id=tokenizer.eos_token_id)
19output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
20output = output.split('```')[2]
21print(output)