code: import os
  from datasets import get_dataset
  from langchain.chat_models import ChatOpenAI
  from langchain.prompts import PromptTemplate
  from langchain.chains import LLMChain
  from langchain.callbacks import get_openai_callback
  from utils.timer import Timer
  import logging
  import numpy as np
  import seaborn as sns
  import matplotlib.pyplot as plt
  
  from config import api_key, load_config
  import wandb
  import langchain
  langchain.verbose = True
  
  # If you don't want your script to sync to the cloud
  os.environ["WANDB_MODE"] = "offline"
  
  CLASSIFIY_PROMPT = """
  You are a text-to-SQL expert able to identify poorly formulated questions in natural language.
  The dataset used is consisting of questions and their corresponding golden SQL queries. You will be given the database schema of the database corresponding to the question.
  Furthermore, you will also be given a hint that provides additional information that is needed to correctly convert the question and interpret the database schema.  
  However, some of the questions in the data are poorly formulated or contain errors. 
  
  Below is a classification scheme for the questions that are to be converted into SQL queries. 
  
  0 = Correct question. May still contain minor errors in language or minor ambiguities that do not affect the interpretation and generation of the SQL query
  1 = Is unclear, ambiguous, unspecific or contain grammatical errors that surely is going to affect the interpretation and generation of the SQL query. The question 
  is unspecific in which columns that are to be returned. The question is not asking for a specific column, but asks generally about a table in the database.
  2 = The question contains minor errors in language or minor ambiguities that might affect the interpretation and generation of the SQL query.
  3 = The question is wrongly formulated when considering the structure of the database schema. The information that the question is asking for is not possible to accurately retrieve from the database.
  
  Here are some examples of questions that would be classified with 1 and an explanation of why:
  
  Example 1: List the customer who made the transaction id 3682978
  Explanation: The question is unspecific in which columns that are to be returned. It asks to list the customers, but does not specify which columns that are to be returned from the client table. 
  
  Example 2: Which district has the largest amount of female clients?
  Explanation: The question is unspecific in which columns that are to be returned. It asks "which district", but does not specify which columns that are to be returned from the district table. 
  
  Example 3: What is the average amount of transactions done in the year of 1998 ?
  Explanation: Is unclear, ambiguous, unspecific or contain grammatical errors that surely is going to affect the interpretation and generation of the SQL query.
  
  Here is an example of a question that would be classified with 2 and an explanation of why:
  
  Example 1: What are the top 5 loans by region names for the month of Mars 1997?
  Explanation: The statement 'top 5' could be ambiguous. It could mean the top 5 loans by amount or the top 5 loans by number of loans.
  
  Here are some examples of questions that would be classified with 3 and an explanation of why:
  
  Example 1: What is the disposition id of the oldest client in the Prague region?
  Explanation: The question is wrongly formulated when considering the structure of the database schema. There can be multiple disposition ids for a client, 
  since a client can have multiple accounts. The question is not asking for a specific disposition id, but asks generally about a client.
  
  
  Here are some examples of questions that would be classified with 0 and an explanation of why:
  
  Example 1: List the id of the customer who made the transaction id : 3682978
  Explanation: Clear and correct question.
  
  Example 2: What is the name of the district that has the largest amount of female clients?
  Explanation: Specific and  correct question.
  
  Example 3: What is the disposition id(s) of the oldest client in the Prague region?
  Explanation: The question is open for disposition ids which is correct when considering the sql-schema.
  
  Example 4: What was the average number of withdrawal transactions conducted by female clients from the Prague region during the year 1998?
  Explanation: Clear and correct question.
  
  Database schema: 
  {database_schema}
  
  Hint:
  {evidence}
  
  Below you will be provided with the correct SQL-query that represents what the questions is trying to ask for.
  
  Gold query: 
  {gold_query}
  
  Please classify the question below according to the classification scheme above, the examples, the hint and the SQL gold query provided.
  Also please assume that all dates, values, names and numbers in the questions are correct. 
  
  Question: 
  {question}
  
  In your answer DO NOT return anything else than the mark as a sole number. Do not return any corresponding text or explanations. 
  """
  
  #1 = Gray area, minor errors that may or may not affect the interpretation and generation of the SQL query.
  
  
  class Classifier():
      total_tokens = 0
      prompt_tokens = 0 
      total_cost = 0
      completion_tokens = 0
      last_call_execution_time = 0
      total_call_execution_time = 0
  
      def __init__(self, llm):        
          self.llm = llm
  
          self.prompt_template = CLASSIFIY_PROMPT
          prompt = PromptTemplate(    
              # input_variables=["question", "database_schema","evidence"],
              input_variables=["question", "database_schema", "evidence", 'gold_query'],
              template=CLASSIFIY_PROMPT,
          )
  
          self.chain = LLMChain(llm=llm, prompt=prompt)
  
  
      def classify_question(self, question, schema, evidence, gold_query):
          with get_openai_callback() as cb:
              with Timer() as t:
                  response = self.chain.run({
                      'question': question,
                      'database_schema': schema,
                      'evidence': evidence,
                      'gold_query': gold_query
                  })
  
              logging.info(f"OpenAI API execution time: {t.elapsed_time:.2f}")
              
              self.last_call_execution_time = t.elapsed_time
              self.total_call_execution_time += t.elapsed_time
              self.total_tokens += cb.total_tokens
              self.prompt_tokens += cb.prompt_tokens
              self.total_cost += cb.total_cost
              self.completion_tokens += cb.completion_tokens
  
              return response
  
  
  accepted_faults = [1, 3]
  
  def main():
      config = load_config("classifier_config.yaml")
  
      wandb.init(
          project=config.project,
          config=config,
          name=config.current_experiment,
          entity=config.entity
      )
  
      artifact = wandb.Artifact('experiment_results', type='dataset')
      table = wandb.Table(columns=["Question", "Classified_quality", "Difficulty"]) ## Är det något mer vi vill ha med här?
      wandb_cm = wandb.Table(columns=['0', '1', '2', '3'])
      metrics_table = wandb.Table(columns=["Class", "Precision", "Recall", "F1 Score", "Accuracy"])
      weighted_avg_table = wandb.Table(columns=["Metric", "Weighted Average"])
      # "Weighted Averages", weighted_averages['precision'], weighted_averages['recall'], weighted_averages['f1'], weighted_averages['accuracy']
  
  
      llm = ChatOpenAI(
          openai_api_key=api_key, 
          model_name=config.llm_settings.model,
          temperature=config.llm_settings.temperature,
          request_timeout=config.llm_settings.request_timeout
      )
  
      dataset = get_dataset("BIRDCorrectedFinancialGoldAnnotated")
      classifier = Classifier(llm)
  
      wandb.config['prompt'] = classifier.prompt_template
  
      no_data_points = dataset.get_number_of_data_points()
  
      tp = 0
      fp = 0
      tn = 0
      fn = 0
      confusion_matrix = np.zeros((4,4))
      annotation_counts = {0: 0, 1: 0, 2: 0, 3: 0}
      
      for i in range(no_data_points):
          data_point = dataset.get_data_point(i)
          evidence = data_point['evidence']
          db_id = data_point['db_id']            
          question = data_point['question']
          gold_query = data_point['SQL']
          difficulty = data_point['difficulty'] if 'difficulty' in data_point else ""
          annotated_question_quality = data_point["annotation"]
          
          sql_schema = dataset.get_schema_and_sample_data(db_id)
  
          classified_quality = classifier.classify_question(question, sql_schema, evidence, gold_query)
          classified_quality = int(classified_quality) if classified_quality.isdigit() else None
  
          print('classified_quality: ',classified_quality)
  
          if classified_quality is not None:
              for annotated_quality in annotated_question_quality:  
                  annotation_counts[annotated_quality] +=1
                  confusion_matrix[annotated_quality][classified_quality] += 1
                  
  
  
      print('confusion matrix:')
      print(confusion_matrix)
      # Converting to integer
      confusion_matrix = np.array(confusion_matrix).astype(int)
      
      print('annotation counts: ',annotation_counts)
      labels = [0, 1, 2, 3] 
      sns.heatmap(confusion_matrix, annot=True, fmt="d", cmap="YlOrRd", xticklabels=labels, yticklabels=labels)
      plt.ylabel('True label')
      plt.xlabel('Predicted label')
      
      plt.savefig(f'{config.current_experiment}_heatmap.png')
  
      wandb.log({"confusion_matrix_heatmap": wandb.Image(f'{config.current_experiment}_heatmap.png')})
  
      metrics = {'precision': 0, 'recall': 0, 'f1': 0, 'accuracy': 0}
      weighted_sums = {'precision': 0, 'recall': 0, 'f1': 0, 'accuracy': 0}
      total_instances = np.sum(confusion_matrix)
  
      for i in range(4):
          row_data = confusion_matrix[i].tolist()
          print('row_data: ', row_data)
          wandb_cm.add_data(*row_data)
          tp = confusion_matrix[i][i]
          fp = sum(confusion_matrix[:, i]) - tp
          fn = sum(confusion_matrix[i, :]) - tp
          tn = np.sum(confusion_matrix) - (tp + fp + fn)
  
          precision = tp / (tp + fp) if (tp + fp) != 0 else 0
          recall = tp / (tp + fn) if (tp + fn) != 0 else 0
          f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) != 0 else 0
          accuracy = (tp + tn) / (tp + tn + fp + fn)
  
          metrics[i] = {'precision': precision, 'recall': recall, 'f1': f1, 'accuracy': accuracy}
          metrics_table.add_data(i, metrics[i]['precision'], metrics[i]['recall'], metrics[i]['f1'], metrics[i]['accuracy'])
  
          class_weight = sum(confusion_matrix[i, :])
          weighted_sums['precision'] += precision * class_weight
          weighted_sums['recall'] += recall * class_weight
          weighted_sums['f1'] += f1 * class_weight
          weighted_sums['accuracy'] += accuracy * class_weight
  
          print('metrics for class ', i, ': ', metrics[i])
  
  
  
  
          # metrics now contains the precision, recall, and F1-score for each category
              # annotated_question_qualities = set(annotated_question_quality)
              # if classified_quality.isdigit() and int(classified_quality) == 1:            
              #     if any(element in annotated_question_qualities for element in accepted_faults):
              #         tp += 1
              #     else:
              #         fp += 1
              # elif classified_quality.isdigit() and int(classified_quality) == 0:
              #     if any(element in annotated_question_qualities for element in accepted_faults):
              #         fn += 1
              #     else:
              #         tn += 1
              
              # precision = tp / (tp + fp)
              # recall = tp / (tp + fn)
              # f1 = 2 * ((precision * recall) / (precision + recall))
              # accuracy = (tp + tn) / (tp + tn + fp + fn)
  
          table.add_data(question, classified_quality, difficulty)
          wandb.log({                      
              "total_tokens": classifier.total_tokens,
              "prompt_tokens": classifier.prompt_tokens,
              "completion_tokens": classifier.completion_tokens,
              "total_cost": classifier.total_cost,
              "openAPI_call_execution_time": classifier.last_call_execution_time,
          }, step=i+1)
      
          print("Predicted quality: ", classified_quality, " Annotated quality: ", " ".join(map(str, annotated_question_quality)))
          
      weighted_averages = {metric: total / total_instances for metric, total in weighted_sums.items()}
  
      print("Weighted Averages:", weighted_averages)
      
      # weighted_avg_table.add_data("Weighted Averages", weighted_averages['precision'], weighted_averages['recall'], weighted_averages['f1'], weighted_averages['accuracy'])
      weighted_avg_table.add_data("Precision", weighted_averages['precision'])
      weighted_avg_table.add_data("Recall", weighted_averages['recall'])
      weighted_avg_table.add_data("F1 Score", weighted_averages['f1'])
      weighted_avg_table.add_data("Accuracy", weighted_averages['accuracy'])
  
  
  
      wandb.run.summary["total_tokens"]                       = classifier.total_tokens
      wandb.run.summary["prompt_tokens"]                      = classifier.prompt_tokens
      wandb.run.summary["completion_tokens"]                  = classifier.completion_tokens
      wandb.run.summary["total_cost"]                         = classifier.total_cost
      wandb.run.summary['total_predicted_execution_time']     = dataset.total_predicted_execution_time
      wandb.run.summary['total_openAPI_execution_time']       = classifier.total_call_execution_time
  
      
      artifact.add(wandb_cm, "ConfusionMatrix_predictions")
      artifact.add(table, "query_results")
      artifact.add(metrics_table, "metrics")
      artifact.add(weighted_avg_table, "weighted_averages_metric_table")
      wandb.log_artifact(artifact)
  
      artifact_code = wandb.Artifact('code', type='code')
      artifact_code.add_file("src/run_classifier.py")
      wandb.log_artifact(artifact_code)
  
      wandb.finish()
  
  
  
  if __name__ == "__main__":
      main()
apis: langchain.prompts.PromptTemplate langchain.callbacks.get_openai_callback langchain.chat_models.ChatOpenAI langchain.chains.LLMChain
extract_api: [((6515, 6552), 'config.load_config', 'load_config', (['"""classifier_config.yaml"""'], {}), "('classifier_config.yaml')\n", (6526, 6552), False, 'from config import api_key, load_config\n'), ((6558, 6666), 'wandb.init', 'wandb.init', ([], {'project': 'config.project', 'config': 'config', 'name': 'config.current_experiment', 'entity': 'config.entity'}), '(project=config.project, config=config, name=config.\n    current_experiment, entity=config.entity)\n', (6568, 6666), False, 'import wandb\n'), ((6716, 6768), 'wandb.Artifact', 'wandb.Artifact', (['"""experiment_results"""'], {'type': '"""dataset"""'}), "('experiment_results', type='dataset')\n", (6730, 6768), False, 'import wandb\n'), ((6781, 6850), 'wandb.Table', 'wandb.Table', ([], {'columns': "['Question', 'Classified_quality', 'Difficulty']"}), "(columns=['Question', 'Classified_quality', 'Difficulty'])\n", (6792, 6850), False, 'import wandb\n'), ((6906, 6947), 'wandb.Table', 'wandb.Table', ([], {'columns': "['0', '1', '2', '3']"}), "(columns=['0', '1', '2', '3'])\n", (6917, 6947), False, 'import wandb\n'), ((6968, 7045), 'wandb.Table', 'wandb.Table', ([], {'columns': "['Class', 'Precision', 'Recall', 'F1 Score', 'Accuracy']"}), "(columns=['Class', 'Precision', 'Recall', 'F1 Score', 'Accuracy'])\n", (6979, 7045), False, 'import wandb\n'), ((7071, 7122), 'wandb.Table', 'wandb.Table', ([], {'columns': "['Metric', 'Weighted Average']"}), "(columns=['Metric', 'Weighted Average'])\n", (7082, 7122), False, 'import wandb\n'), ((7278, 7457), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'openai_api_key': 'api_key', 'model_name': 'config.llm_settings.model', 'temperature': 'config.llm_settings.temperature', 'request_timeout': 'config.llm_settings.request_timeout'}), '(openai_api_key=api_key, model_name=config.llm_settings.model,\n    temperature=config.llm_settings.temperature, request_timeout=config.\n    llm_settings.request_timeout)\n', (7288, 7457), False, 'from langchain.chat_models import ChatOpenAI\n'), ((7503, 7553), 'datasets.get_dataset', 'get_dataset', (['"""BIRDCorrectedFinancialGoldAnnotated"""'], {}), "('BIRDCorrectedFinancialGoldAnnotated')\n", (7514, 7553), False, 'from datasets import get_dataset\n'), ((7770, 7786), 'numpy.zeros', 'np.zeros', (['(4, 4)'], {}), '((4, 4))\n', (7778, 7786), True, 'import numpy as np\n'), ((9068, 9177), 'seaborn.heatmap', 'sns.heatmap', (['confusion_matrix'], {'annot': '(True)', 'fmt': '"""d"""', 'cmap': '"""YlOrRd"""', 'xticklabels': 'labels', 'yticklabels': 'labels'}), "(confusion_matrix, annot=True, fmt='d', cmap='YlOrRd',\n    xticklabels=labels, yticklabels=labels)\n", (9079, 9177), True, 'import seaborn as sns\n'), ((9178, 9202), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""True label"""'], {}), "('True label')\n", (9188, 9202), True, 'import matplotlib.pyplot as plt\n'), ((9207, 9236), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Predicted label"""'], {}), "('Predicted label')\n", (9217, 9236), True, 'import matplotlib.pyplot as plt\n'), ((9246, 9301), 'matplotlib.pyplot.savefig', 'plt.savefig', (['f"""{config.current_experiment}_heatmap.png"""'], {}), "(f'{config.current_experiment}_heatmap.png')\n", (9257, 9301), True, 'import matplotlib.pyplot as plt\n'), ((9569, 9593), 'numpy.sum', 'np.sum', (['confusion_matrix'], {}), '(confusion_matrix)\n', (9575, 9593), True, 'import numpy as np\n'), ((13640, 13668), 'wandb.log_artifact', 'wandb.log_artifact', (['artifact'], {}), '(artifact)\n', (13658, 13668), False, 'import wandb\n'), ((13690, 13725), 'wandb.Artifact', 'wandb.Artifact', (['"""code"""'], {'type': '"""code"""'}), "('code', type='code')\n", (13704, 13725), False, 'import wandb\n'), ((13782, 13815), 'wandb.log_artifact', 'wandb.log_artifact', (['artifact_code'], {}), '(artifact_code)\n', (13800, 13815), False, 'import wandb\n'), ((13821, 13835), 'wandb.finish', 'wandb.finish', ([], {}), '()\n', (13833, 13835), False, 'import wandb\n'), ((5346, 5466), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['question', 'database_schema', 'evidence', 'gold_query']", 'template': 'CLASSIFIY_PROMPT'}), "(input_variables=['question', 'database_schema', 'evidence',\n    'gold_query'], template=CLASSIFIY_PROMPT)\n", (5360, 5466), False, 'from langchain.prompts import PromptTemplate\n'), ((5598, 5630), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt'}), '(llm=llm, prompt=prompt)\n', (5606, 5630), False, 'from langchain.chains import LLMChain\n'), ((11746, 12026), 'wandb.log', 'wandb.log', (["{'total_tokens': classifier.total_tokens, 'prompt_tokens': classifier.\n    prompt_tokens, 'completion_tokens': classifier.completion_tokens,\n    'total_cost': classifier.total_cost, 'openAPI_call_execution_time':\n    classifier.last_call_execution_time}"], {'step': '(i + 1)'}), "({'total_tokens': classifier.total_tokens, 'prompt_tokens':\n    classifier.prompt_tokens, 'completion_tokens': classifier.\n    completion_tokens, 'total_cost': classifier.total_cost,\n    'openAPI_call_execution_time': classifier.last_call_execution_time},\n    step=i + 1)\n", (11755, 12026), False, 'import wandb\n'), ((5719, 5740), 'langchain.callbacks.get_openai_callback', 'get_openai_callback', ([], {}), '()\n', (5738, 5740), False, 'from langchain.callbacks import get_openai_callback\n'), ((6031, 6095), 'logging.info', 'logging.info', (['f"""OpenAI API execution time: {t.elapsed_time:.2f}"""'], {}), "(f'OpenAI API execution time: {t.elapsed_time:.2f}')\n", (6043, 6095), False, 'import logging\n'), ((8942, 8968), 'numpy.array', 'np.array', (['confusion_matrix'], {}), '(confusion_matrix)\n', (8950, 8968), True, 'import numpy as np\n'), ((9346, 9401), 'wandb.Image', 'wandb.Image', (['f"""{config.current_experiment}_heatmap.png"""'], {}), "(f'{config.current_experiment}_heatmap.png')\n", (9357, 9401), False, 'import wandb\n'), ((9882, 9906), 'numpy.sum', 'np.sum', (['confusion_matrix'], {}), '(confusion_matrix)\n', (9888, 9906), True, 'import numpy as np\n'), ((5765, 5772), 'utils.timer.Timer', 'Timer', ([], {}), '()\n', (5770, 5772), False, 'from utils.timer import Timer\n')]