core/management/commands/text2sql_eval.py

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}")