Views
No views yet
1import torch
2from transformers import AutoTokenizer, pipeline
3from datasets import load_dataset
4from peft import AutoPeftModelForCausalLM
5from random import randint
6
7peft_model_id = "delayedkarma/mistral-7b-text-to-sql"
8
9# Load Model with PEFT adapter
10model = AutoPeftModelForCausalLM.from_pretrained(
11 peft_model_id,
12 device_map="auto",
13 torch_dtype=torch.float16
14)
15tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
16# load into pipeline
17pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
18
19# Load dataset and Convert dataset to OAI messages
20system_message = """You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA.
21SCHEMA:
22{schema}"""
23
24def create_conversation(sample):
25 return {
26 "messages": [
27 {"role": "system", "content": system_message.format(schema=sample["context"])},
28 {"role": "user", "content": sample["question"]},
29 {"role": "assistant", "content": sample["answer"]}
30 ]
31 }
32
33# Load dataset from the hub
34dataset = load_dataset("b-mc2/sql-create-context", split="train")
35dataset = dataset.shuffle().select(range(100))
36
37# Convert dataset to OAI messages
38dataset = dataset.map(create_conversation, remove_columns=dataset.features, batched=False)
39
40dataset = dataset.train_test_split(test_size=20/100)
41
42# Evaluate
43eval_dataset = dataset['test']
44rand_idx = randint(0, len(eval_dataset))
45
46# Test on sample
47prompt = pipe.tokenizer.apply_chat_template(eval_dataset[rand_idx]["messages"][:2], tokenize=False, add_generation_prompt=True)
48outputs = pipe(prompt, max_new_tokens=256, do_sample=False, temperature=0.1, top_k=50, top_p=0.1, eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)
49
50print(f"Query:\n{eval_dataset[rand_idx]['messages'][1]['content']}")
51print(f"Original Answer:\n{eval_dataset[rand_idx]['messages'][2]['content']}")
52print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip()}")
53