run.py

import fire
import pandas as pd
import ollama
import re
from tqdm import tqdm
from pathlib import Path
#MODELS = ['codegemma:7b-instruct-q4_0']
MODELS = ['codegemma:7b-instruct-fp16', 'codellama:34b-instruct-q4_0',
          'deepseek-coder:33b-instruct-q4_0', 'codestral:22b-v0.1-q4_0', 
          'qwen2.5-coder:32b-instruct-q4_K_M', 'deepseek-coder-v2:16b-lite-instruct-q4_0']

EXCEPTION_TYPES = {
  'python': [
      'NameError',
      'TypeError',
      'IndexError',
      'ValueError',
      'AttributeError',
      'ZeroDivisionError',
      'KeyError',
      'SyntaxError',
      'UnboundLocalError',
      'OverflowError',
      'RecursionError'
    ],

  'c_sharp': [
    'System.IndexOutOfRangeException',
    'System.OverflowException',
    'System.ArgumentOutOfRangeException',
    'System.DivideByZeroException',
    'System.FormatException',
    'System.NullReferenceException',
    'System.InvalidOperationException',
    'System.Collections.Generic.KeyNotFoundException',
    'System.ArgumentNullException'
    ],

  'java': [
    'java.lang.ArrayIndexOutOfBoundsException',
    'java.util.InputMismatchException',
    'java.util.NoSuchElementException',
    'java.lang.StringIndexOutOfBoundsException',
    'java.lang.NumberFormatException',
    'java.lang.NullPointerException',
    'java.lang.ArithmeticException',
    'java.lang.IndexOutOfBoundsException',
    'java.lang.StackOverflowError',
    'java.lang.NegativeArraySizeException'
  ],

  'ruby': [
    'NoMethodError',
    'NameError',
    'TypeError',
    'ArgumentError',
    'ZeroDivisionError',
    'FloatDomainError',
    'RangeError',
  ]

}

LANG_NAMES = {
    'python': 'Python',
    'java': 'Java',
    'c_sharp': 'C#',
    'ruby': 'Ruby'
}

RAISE = ['raise', 'raises', 'raised']
THROW = ['throw', 'throws', 'thrown']

RAISE_OR_THROW = {
    'python': RAISE,
    'ruby': RAISE,
    'java': THROW,
    'c_sharp': THROW,
}

def generate_prompt_buggy(row, language, explain=False):
    lang_name = LANG_NAMES[language]
    exceptions = EXCEPTION_TYPES[language]
    throw, throws, thrown = RAISE_OR_THROW[language]
    if not explain:
      explain_instruction = f"Only output your answer. Do not output explanations."
    else:
      explain_instruction = f"Please explain your decision/reasoning. Give your answer at the end"
    return f"""
You are given a {lang_name} program that either does not {throw} an exception or {throws} one of the following exceptions:
   {', '.join(exceptions)}
when run with the provided program input.
Please find and output the name of the exception (or output \"no_exception\" if no exception is {thrown}).
Give your answer in the following format: "Answer: <exception_name>" or "Answer: no_exception" (if no exception is {thrown}).
{explain_instruction}   

Code:
----------------
{row.code}
----------------

Program Input:
---------------
{row.input}
---------------
"""

def generate_prompt(row, language, format='buggy', explain=False):
    if format == 'buggy':
        return generate_prompt_buggy(row, language, explain=explain)
    else:
        raise ValueError(f'invalid format "{format}"')

def extract_exception(language, output):
  exception_types_regex = '|'.join(EXCEPTION_TYPES[language])
  matches = re.findall(rf"{exception_types_regex}|no_exception", output)
  if matches:
    return matches[-1]
  else:
    print("No answer found")
    return None

def main(data_file, *, output_dir, explain=False):
    df = pd.read_json(data_file, orient='record', lines=True)
    df.exception_type = df.exception_type.fillna(value='no_exception')

    for l, g in df.groupby('language'):
      print(sorted(EXCEPTION_TYPES[l] + ['no_exception']))
      print(sorted(g.exception_type.unique().tolist()))
      assert sorted(EXCEPTION_TYPES[l] + ['no_exception']) == sorted(g.exception_type.unique().tolist())

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    for model in MODELS:
        output_rows = []
        tuples = list(df.itertuples(index=False))
        for i, row in enumerate(tqdm(tuples, desc=model)):
            prompt = generate_prompt(row, row.language, 'buggy', explain=explain)
            try:
              if explain:
                response = ollama.generate(model=model, prompt=prompt, options=dict(num_predict=1024))
              else:
                response = ollama.generate(model=model, prompt=prompt)
              output = response.response.strip()
              predicted_exception = extract_exception(row.language, output)
            except ollama.ResponseError as e:
              print('Error:', e.error)
              predicted_exception = 'ERROR'

            if i % 100 == 0:
                print(prompt)
                print("------------------")
                print(output)
                print("------------------")
                print(predicted_exception)
                print("\n\n\n\n")
            output_rows.append({'model': model, 'predicted_exception': predicted_exception,
                                'actual_exception': row.exception_type,
                                'bug_id': row.bug_id,
                                'output': output,
                                'language': row.language})

        output_df = pd.DataFrame(output_rows)
        output_df['correct'] = output_df.predicted_exception == output_df.actual_exception
        print(output_df.groupby(['model', 'language'])['correct'].mean())
        output_df.to_csv(output_path / f"{model}.csv.gz", index=False)

if __name__ == '__main__':
    fire.Fire(main)