1from transformers import (
2 AutoModelForCausalLM,
3 AutoTokenizer,
4 pipeline
5)
6import torch
7
8model_name = 'bugdaryan/MistralSQL-7b'
9
10model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto')
11tokenizer = AutoTokenizer.from_pretrained(model_name)
12
13pipe = pipeline('text-generation', model=model, tokenizer=tokenizer)
14
15table = "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 );"
16
17question = 'Find the salesperson who made the most sales.'
18
19prompt = 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}: ``` "
20
21ans = pipe(prompt, max_new_tokens=100)
22print(ans[0]['generated_text'].split('```')[2])