File size: 1,859 Bytes
38171fa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import logging
import time
import csv
from django.core.management.base import BaseCommand
from core.text2sql.handler import QueryDataHandler
from core.text2sql.prompt import get_prompt
from core.text2sql.eval_queries import queries
logger = logging.getLogger(__name__)
class Command(BaseCommand):
help = "Morningstar API to save the JSON response to a file which contains secIds with other details"
def handle(self, *args, **options) -> None:
t1 = time.perf_counter()
q=[]
count =1
for query in queries[26:]:
print("count: ", query["Query Number"])
prompt = get_prompt(query["Query Description"])
logger.info(f"Prompt: {prompt}")
generated_query, data = QueryDataHandler().get_data_from_query(prompt)
print(f"Description: {query['Query Description']}, Query: {query.get('SQL Statement')}, Generated: {generated_query} ")
q.append({
"Query Number": query["Query Number"],
"Complexity Level": query["Complexity Level"],
"Description": query["Query Description"],
"Query": query.get("SQL Statement", "-"),
"Generated": generated_query,
})
count+=1
time.sleep(1)
csv_file_path = 'queries_data.csv'
# Writing data to CSV
with open(csv_file_path, 'w', newline='', encoding='utf-8') as csv_file:
fieldnames = q[0].keys()
print(fieldnames)
writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
# Write the header
writer.writeheader()
# Write the data
writer.writerows(q)
print(f'Data has been written to {csv_file_path}.')
self.stdout.write(f"Time taken for evaluation: {time.perf_counter() - t1}")
|