Views
No views yet
Chirayu TripathiMicrosoft]microsoft/phi-2prompt_template = f"""<s>
Task Description:
Your task is to create a MongoDB query that accurately fulfills the provided Instruct while strictly adhering to the given MongoDB schema. Ensure that the query solely relies on keys and columns present in the schema. Minimize the usage of lookup operations wherever feasible to enhance query efficiency.
MongoDB Schema:
{db_schema}
### Instruct:
{text}
### Output:
"""1from transformers import (
2 AutoTokenizer,
3 AutoModelForCausalLM,
4 BitsAndBytesConfig,
5)
6import torch
7from peft import PeftModel
8
9db_schema = '''{
10 "collections": [
11 {
12 "name": "shipwrecks",
13 "indexes": [
14 {
15 "key": {
16 "_id": 1
17 }
18 },
19 {
20 "key": {
21 "feature_type": 1
22 }
23 },
24 {
25 "key": {
26 "chart": 1
27 }
28 },
29 {
30 "key": {
31 "latdec": 1,
32 "londec": 1
33 }
34 }
35 ],
36 "uniqueIndexes": [],
37 "document": {
38 "properties": {
39 "_id": {
40 "bsonType": "string"
41 },
42 "recrd": {
43 "bsonType": "string"
44 },
45 "vesslterms": {
46 "bsonType": "string"
47 },
48 "feature_type": {
49 "bsonType": "string"
50 },
51 "chart": {
52 "bsonType": "string"
53 },
54 "latdec": {
55 "bsonType": "double"
56 },
57 "londec": {
58 "bsonType": "double"
59 },
60 "gp_quality": {
61 "bsonType": "string"
62 },
63 "depth": {
64 "bsonType": "string"
65 },
66 "sounding_type": {
67 "bsonType": "string"
68 },
69 "history": {
70 "bsonType": "string"
71 },
72 "quasou": {
73 "bsonType": "string"
74 },
75 "watlev": {
76 "bsonType": "string"
77 },
78 "coordinates": {
79 "bsonType": "array",
80 "items": {
81 "bsonType": "double"
82 }
83 }
84 }
85 }
86 }
87 ],
88 "version": 1
89}'''
90
91text = ''''Find the count of shipwrecks for each unique combination of "latdec" and "longdec"'''
92prompt = f"""<s>
93 Task Description:
94 Your task is to create a MongoDB query that accurately fulfills the provided Instruct while strictly adhering to the given MongoDB schema. Ensure that the query solely relies on keys and columns present in the schema. Minimize the usage of lookup operations wherever feasible to enhance query efficiency.
95
96 MongoDB Schema:
97 {db_schema}
98
99 ### Instruct:
100 {text}
101
102 ### Output:
103 """
104
105device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
106base_model_id = "microsoft/phi-2"
107tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_fast=True)
108compute_dtype = getattr(torch, "float16")
109bnb_config = BitsAndBytesConfig(
110 load_in_4bit=True,
111 bnb_4bit_quant_type="nf4",
112 bnb_4bit_compute_dtype=compute_dtype,
113 bnb_4bit_use_double_quant=True,
114)
115model = AutoModelForCausalLM.from_pretrained(
116 base_model_id,
117 trust_remote_code=True,
118 quantization_config=bnb_config,
119 revision="refs/pr/23",
120 device_map={"": 0},
121 torch_dtype="auto",
122 flash_attn=True,
123 flash_rotary=True,
124 fused_dense=True,
125)
126adapter = 'Chirayu/phi-2-mongodb'
127
128model = PeftModel.from_pretrained(model, adapter).to(device)
129model_inputs = tokenizer(prompt, return_tensors="pt").to(device)
130output = model.generate(
131 **model_inputs,
132 max_length=1024,
133 no_repeat_ngram_size=10,
134 repetition_penalty=1.02,
135 pad_token_id=tokenizer.eos_token_id,
136 eos_token_id=tokenizer.eos_token_id,
137)[0]
138
139prompt_length = model_inputs['input_ids'].shape[1]
140query = tokenizer.decode(output[prompt_length:], skip_special_tokens=False)
141try:
142 stop_idx = query.index("</s>")
143except Exception as e:
144 print(e)
145 stop_idx = len(query)
146print(query[: stop_idx].strip())