Merge pull request #22 from vadim0x60/add/ollama-python-package

Add/ollama python package

Author: Anastasiia-Grishina

README.md

@@ -44,12 +44,12 @@ set to your OpenAI API access token.
 
 #### Set up Ollama
 
-Run [Ollama](https://ollama.ai/) with CodeLlama or [another model](https://ollama.ai/library) locally 
+Run [Ollama](https://ollama.ai/) with Llama 3-8B or [another model](https://ollama.ai/library) locally 
 or on a server. 
 In the latter case, start the Ollama server with the following commands and note the `URL:PORT` pair:
 ```
 OLLAMA_HOST=URL:PORT ollama serve &
-OLLAMA_HOST=URL:PORT ollama pull codellama:34b-instruct &
+OLLAMA_HOST=URL:PORT ollama pull llama3 &
 ```
 
 Example `.config` file layout:
@@ -75,41 +75,48 @@ export WANDB_DIR=...
 
 If you're using [Slurm](https://slurm.schedmd.com/), write a `run.sh` file with `python benchmark.py` 
 and run it with `sbatch run.sh --array=1-500`.
-If not, run `TASK_ID=n python benchmark.py` to re-run one of our experiments exactly, or set the parameters yourself:
+If not, run `TASK_ID=n python benchmark.py` to re-run one of our experiments exactly, 
+or set the parameters yourself as below.
 
 For example, for basement problem in PSB2, run SEIDR without lexicase selection as follows:
 ```
 python3 benchmark.py \
-    --task_id 202 \
-    --problem basement \
-    --language C++ \
+    --task_id 0 \
+    --problem bowling \
+    --language Python \
+    --branching_factor 2 \
     --max_programs 100 \
     --drafts_per_prompt 2 \
     --explanations_per_program 2 \
     --repairs_per_explanation 2 \
     --beam_width 2 \
     --log INFO \
     --lexicase_selection False \
-    --dataset psb2 \
-    --model_name gpt-3.5-turbo
+    --dataset humaneval \
+    --model_name gpt-3.5-turbo \
+    --valid_examples 50 \
+    --experiment_id 0
 ```
 
-To run an example with SEIDR with CodeLlama served by Ollama at `URL:PORT`, run the following:
+To run an example with SEIDR with Llama 3 served by Ollama at `URL:PORT` on HumanEval with lexicase, run the following:
 ```
-python3 benchmark.py \
-    --task_id 2202 \
-    --problem basement \
-    --language C++ \
+python3 benchmark_humaneval.py \
+    --task_id 0 \
+    --problem Python/0 \
+    --language Python \
+    --branching_factor 2 \
     --max_programs 100 \
     --drafts_per_prompt 2 \
     --explanations_per_program 2 \
     --repairs_per_explanation 2 \
     --beam_width 2 \
     --log INFO \
-    --lexicase_selection False \
-    --dataset psb2 \
-    --model_name codellama:34b-instruct \
+    --lexicase_selection True \
+    --dataset humaneval \
+    --model_name llama3 \
+    --experiment_id 0 \
     --ollama_url "http://URL:PORT"
+
 ```
 
-Example Slurm scripts are stored in `example_scripts/` and tables with hyperparameters in `/config`
+Example Slurm scripts are stored in `scripts/` and tables with hyperparameters in `/config`

benchmark.py

@@ -20,59 +20,68 @@
 
 logger = logging.getLogger(__name__)
 
-DATA_PATH = os.environ.get('DATA_PATH') or 'psb2'
+DATA_PATH = os.environ.get("DATA_PATH") or "psb2"
 
 task_descriptions = []
-with open('psb2-meta/tasks.txt') as f:
-    task_descriptions = {name.strip(): description.strip()
-                         for name, description in chunked(f.readlines(), 2)}
+with open("psb2-meta/tasks.txt") as f:
+    task_descriptions = {
+        name.strip(): description.strip()
+        for name, description in chunked(f.readlines(), 2)
+    }
 
-debug_templates = [line.split('\t')
-                   for line in get_template('prompts.txt').splitlines()]
-debug_templates = {int(ix.strip()): prompt.strip()
-                   for ix, prompt in debug_templates}
+debug_templates = [
+    line.split("\t") for line in get_template("prompts.txt").splitlines()
+]
+debug_templates = {int(ix.strip()): prompt.strip() for ix, prompt in debug_templates}
 
 
 def title2kebabcase(title: str) -> str:
     """Replace spaces with hyphens"""
-    return '-'.join(word.lower() for word in title.split(' '))
+    return "-".join(word.lower() for word in title.split(" "))
 
 
-pushgp_success_rates = pd.read_csv('psb2-meta/results.tsv',
-                                   sep='\t', index_col=['Problem'])
-pushgp_success_rates = pushgp_success_rates['Succ.'].rename(title2kebabcase)
+pushgp_success_rates = pd.read_csv(
+    "psb2-meta/results.tsv", sep="\t", index_col=["Problem"]
+)
+pushgp_success_rates = pushgp_success_rates["Succ."].rename(title2kebabcase)
 
 
 def is_already_solved(
-        solutions_logger: FileLogger,
-        test_data: Tuple[List[str] | str, List[str] | str],
-        language: Language) -> Program | bool:
+    solutions_logger: FileLogger,
+    test_data: Tuple[List[str] | str, List[str] | str],
+    language: Language,
+) -> Program | bool:
     """Checks if the currently logged solution passes all tests in `test_data`.
     Returns False if a Program class instance cannot be created"""
     try:
-        return Program(workdir=solutions_logger.dir,
-                       name=solutions_logger.filename,
-                       language=language).test(test_data)
+        return Program(
+            workdir=solutions_logger.dir,
+            name=solutions_logger.filename,
+            language=language,
+        ).test(test_data)
     except FileNotFoundError:
         return False
 
 
-def run_benchmark(problem: str = 'fizz-buzz',
-                  language: str = 'C++',
-                  max_programs: int = 1000,
-                  drafts_per_prompt: int = 10,
-                  explanations_per_program: int = 10,
-                  repairs_per_explanation: int = 2,
-                  beam_width: int = 100,
-                  seed: int = 42,
-                  valid_examples: int = 100,
-                  test_examples: int = 2000,
-                  prompt_examples: int = 5,
-                  log: str = 'ERROR',
-                  model_name: str = 'gpt-3.5-turbo',
-                  lexicase_selection: bool = False,
-                  ollama_url: Optional[str] = "http://localhost:11434",
-                  **kwargs):
+def run_benchmark(
+    problem: str = "fizz-buzz",
+    language: str = "C++",
+    max_programs: int = 1000,
+    drafts_per_prompt: int = 10,
+    explanations_per_program: int = 10,
+    repairs_per_explanation: int = 2,
+    beam_width: int = 100,
+    seed: int = 42,
+    valid_examples: int = 100,
+    test_examples: int = 2000,
+    prompt_examples: int = 5,
+    log: str = "ERROR",
+    model_name: str = "gpt-3.5-turbo",
+    lexicase_selection: bool = False,
+    ollama_url: Optional[str] = "http://localhost:11434",
+    experiment_id: int = 0,
+    **kwargs,
+):
     """Generate and repair programs in PSB2
 
     Parameters
@@ -115,74 +124,91 @@ def run_benchmark(problem: str = 'fizz-buzz',
         link to the ollama cluster, default is localhost
     """
     # Setup logging
-    Path('logs').mkdir(exist_ok=True)
-    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
-                        datefmt='%m/%d/%Y %H:%M:%S', level=log.upper())
-    logging.info('logging info')
+    Path("logs").mkdir(exist_ok=True)
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=log.upper(),
+    )
+    logging.info("logging info")
     baseline = pushgp_success_rates[problem]
 
     config = {
-        'slurm_job_id': os.environ.get('SLURM_JOB_ID'),
-        'slurm_task_pid': os.environ.get('SLURM_TASK_PID'),
-        'slurm_array_task_id': os.environ.get('SLURM_ARRAY_TASK_ID'),
-        'slurm_array_job_id': os.environ.get('SLURM_ARRAY_JOB_ID'),
-        'task_id': os.environ.get('TASK_ID'),
+        "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
+        "slurm_task_pid": os.environ.get("SLURM_TASK_PID"),
+        "slurm_array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID"),
+        "slurm_array_job_id": os.environ.get("SLURM_ARRAY_JOB_ID"),
+        "task_id": os.environ.get("TASK_ID"),
         **kwargs,
-        **locals()
+        **locals(),
     }
 
-    del config['kwargs']
-    model_name_tag = model_name.replace(':', '_')
+    del config["kwargs"]
+    model_name_tag = model_name.replace(":", "_")
     run = wandb.init(
-        entity=os.environ.get('WANDB_ENTITY'),
-        project=f'seidr-telo-psb2-{model_name_tag}',
-        dir=os.environ.get('WANDB_DIR'),
-        config=config)
-    logger.info(f'Run config {run.config}, W&B: {run.url}')
+        entity=os.environ.get("WANDB_ENTITY"),
+        project=f"seidr-telo-psb2-{model_name_tag}-run{experiment_id}",
+        dir=os.environ.get("WANDB_DIR"),
+        config=config,
+    )
+    logger.info(f"Run config {run.config}, W&B: {run.url}")
 
     language = language_(language)
 
-    commit_msg_template = get_template('commit.txt').format(
-        problem=problem,
-        wandb_url=run.url)
+    commit_msg_template = get_template("commit.txt").format(
+        problem=problem, wandb_url=run.url
+    )
 
-    lexicase_tag = '_lexicase' if lexicase_selection else ""
-    attempts_branch = f'psb_{model_name_tag}_{drafts_per_prompt}x{explanations_per_program}x{repairs_per_explanation}{lexicase_tag}_dev'
-    solutions_branch = f'psb_{model_name_tag}_{drafts_per_prompt}x{explanations_per_program}x{repairs_per_explanation}{lexicase_tag}'
+    lexicase_tag = "_lexicase" if lexicase_selection else ""
+    attempts_branch = f"psb_{model_name_tag}_{drafts_per_prompt}x{explanations_per_program}x{repairs_per_explanation}{lexicase_tag}_run{experiment_id}_dev"
+    solutions_branch = f"psb_{model_name_tag}_{drafts_per_prompt}x{explanations_per_program}x{repairs_per_explanation}{lexicase_tag}_run{experiment_id}"
 
-    attempts_logger = FileLogger(branch=attempts_branch,
-                                 filename=language.source.format(name=problem),
-                                 commit_msg_template=commit_msg_template)
-    solutions_logger = FileLogger(branch=solutions_branch,
-                                  filename=language.source.format(name=problem),
-                                  commit_msg_template=commit_msg_template)
+    attempts_logger = FileLogger(
+        branch=attempts_branch,
+        filename=language.source.format(name=problem),
+        commit_msg_template=commit_msg_template,
+    )
+    solutions_logger = FileLogger(
+        branch=solutions_branch,
+        filename=language.source.format(name=problem),
+        commit_msg_template=commit_msg_template,
+    )
 
     description = task_descriptions[problem]
 
     # ensure that the same I/O pairs are fetched for every experiment
     random.seed(seed)
 
     train_data, test_data = psb2.fetch_examples(
-        DATA_PATH, problem, max(valid_examples, prompt_examples),
-        test_examples, format='competitive')
+        DATA_PATH,
+        problem,
+        max(valid_examples, prompt_examples),
+        test_examples,
+        format="competitive",
+    )
     prompt_data = train_data[:prompt_examples]
     valid_data = train_data[:valid_examples]
 
     if is_already_solved(solutions_logger, test_data, language):
-        logging.info(f'{problem} is already solved, shutting down')
+        logging.info(f"{problem} is already solved, shutting down")
         return
 
     call_count = 0
 
     def log_llm_call(**kwargs):
         """Update and log the number of LLM calls"""
         nonlocal call_count
-        wandb.log({'llm_calls': call_count})
+        wandb.log({"llm_calls": call_count})
         call_count += 1
 
     critics = [
-        lambda code: IOMatch(code=code, language=language, input=inp, output=out,
-                             task_description=description)
+        lambda code: IOMatch(
+            code=code,
+            language=language,
+            input=inp,
+            output=out,
+            task_description=description,
+        )
         for inp, out in valid_data
     ]
     prompt = initial_prompt(description, prompt_data)
@@ -204,31 +230,36 @@ def log_llm_call(**kwargs):
         log_solution=solutions_logger,
         log_llm_call=log_llm_call,
         max_programs=max_programs,
-        ollama_url=ollama_url
+        ollama_url=ollama_url,
     )
 
     solution = seidr.develop(start_code=start_code)
 
-    logging.info('Development done. Testing...')
+    logging.info("Development done. Testing...")
 
     test_evals = [
-        IOMatch(solution,
-                language=language,
-                input=inp, output=out,
-                task_description=description)
-        for inp, out in test_data]
+        IOMatch(
+            solution,
+            language=language,
+            input=inp,
+            output=out,
+            task_description=description,
+        )
+        for inp, out in test_data
+    ]
     avg_score = sum(e.score() for e in test_evals) / len(test_evals)
     test_pass_rate = sum(e.check() for e in test_evals) / len(test_evals)
 
-    logging.info(f'\nTest pass rate on test: {test_pass_rate}\nTest avg score on test: {avg_score}')
+    logging.info(
+        f"\nTest pass rate on test: {test_pass_rate}\nTest avg score on test: {avg_score}"
+    )
 
-    run.log({'test_avg_score': avg_score,
-               'test_pass_rate': test_pass_rate})
+    run.log({"test_avg_score": avg_score, "test_pass_rate": test_pass_rate})
     # run.finish()
     wandb.finish()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     try:
         Fire(run_benchmark)
     except: