Views
No views yet

1# Example code to load and use the model
2from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3
4model_name = "OGSQL-Mistral7B"
5tokenizer = AutoTokenizer.from_pretrained(model_name)
6model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
7
8def generate_sql(query):
9 inputs = tokenizer.encode(query, return_tensors="pt")
10 outputs = model.generate(inputs)
11 return tokenizer.decode(outputs[0], skip_special_tokens=True)
12
13# Example use
14query = """
15using this context:
16-- Create Customers Table
17CREATE TABLE Customers (
18 customer_id INTEGER PRIMARY KEY,
19 name TEXT NOT NULL,
20 email TEXT,
21 join_date DATE
22);
23
24-- Create Products Table
25CREATE TABLE Products (
26 product_id INTEGER PRIMARY KEY,
27 name TEXT NOT NULL,
28 price DECIMAL(10, 2)
29);
30
31-- Create Orders Table
32CREATE TABLE Orders (
33 order_id INTEGER PRIMARY KEY,
34 customer_id INTEGER,
35 product_id INTEGER,
36 order_date DATE,
37 quantity INTEGER,
38 total_price DECIMAL(10, 2),
39 FOREIGN KEY (customer_id) REFERENCES Customers(customer_id),
40 FOREIGN KEY (product_id) REFERENCES Products(product_id)
41);
42
43show me all the orders from last month , sort by date
44
45
46"""
47print(generate_sql(query))
48