# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
# default_exp wandb
# -

# # Weights & Biases Sweeps
#
# > Weights & Biases Sweeps are used to automate hyperparameter optimization and explore the space of possible models.

#export
from fastcore.script import *
from tsai.imports import *
from tsai.utils import *
from tsai.export import *


# +
# # hide
# import os
# from tsai.imports import *
# from tsai.utils import *
# from fastcore.script import *
# from fastcore.xtras import *

# @call_parse
# def run_sweep(
#     sweep:     Param("Path to YAML file with the sweep config", str) = None,
#     program:   Param("Path to Python training script", str) = None,
#     launch:    Param("Launch wanbd agent.", store_false) = True,
#     count:     Param("Number of runs to execute", int) = None,
#     entity:    Param("username or team name where you're sending runs", str) = None,
#     project:   Param("The name of the project where you're sending the new run.", str) = None,
#     sweep_id:  Param("Sweep ID. This option omits `sweep`", str) = None,
#     relogin:   Param("Relogin to wandb.", store_true) = False,
#     login_key: Param("Login key for wandb", str) = None,
# ):

#     # import wandb
#     try:
#         import wandb
#     except ImportError:
#         raise ImportError('You need to install wandb to run sweeps!')

#     # Login to W&B
#     if relogin:
#         wandb.login(relogin=True)
#     elif login_key:
#         wandb.login(key=login_key)

#     # Sweep id
#     if not sweep_id:
#         # Load the sweep config
#         assert os.path.isfile(sweep), f"can't find file {sweep}"
#         if isinstance(sweep, str):
#             sweep = yaml2dict(sweep)
#         if program is None:
#             program = sweep["program"]
#         # Initialize the sweep
#         print('Initializing sweep...')
#         sweep_id = wandb.sweep(sweep=sweep, entity=entity, project=project)
#         print('...sweep initialized')

#     # Load your training script
#     print('Loading training script...')
#     assert program is not None, "you need to pass either a sweep or program path"
#     while True: 
#         if program[0] in "/ .": program = program.split(program[0], 1)[1]
#         else: break
#     if '/' in program and program.rsplit('/', 1)[0] not in sys.path: sys.path.append(program.rsplit('/', 1)[0])
# #     assert os.path.isfile(program), f"can't find file program = {program}"
#     train_script, file_path = import_file_as_module(program, True)
#     assert hasattr(train_script, "train")
#     train_fn = train_script.train
#     print('...training script loaded')

#     # Launch agent
#     if launch:
#         print('\nRun additional sweep agents with:\n')
#     else:
#         print('\nRun sweep agent with:\n')
#     print('    from a notebook:')
#     print('        import wandb')
#     print(f'        from {file_path} import train')
#     print(f"        wandb.agent('{sweep_id}', function=train, count=None)\n")
#     print('    from a terminal:')
#     print(
#         f"        wandb agent {os.environ['WANDB_ENTITY']}/{os.environ['WANDB_PROJECT']}/{sweep_id}\n")
#     if launch:
#         print('Running agent...')
#         try: 
#             wandb.agent(sweep_id, function=train_fn, count=count)
#         except KeyboardInterrupt:
#             pass

# +
# hide
# @call_parse
# def run_sweep(path:      Param('Path to file with the sweep script file', str) = None,
#               entity:    Param("username or team name where you're sending runs", str) = None,
#               project:   Param("The name of the project where you're sending the new run.", str) = None,
#               count:     Param('Number of runs to execute', int) = None,
#               sweep_id:  Param('Sweep ID. This option omits `sweep`', str) = None,
#               launch:    Param("Launch wanbd agent.", store_false) = True,
#               relogin:   Param('Relogin to wandb.', store_true) = False,
#               login_key: Param('Login key for wandb', str) = None,
#               ):

#     try: import wandb
#     except ImportError: raise ImportError('You need to install wandb to run sweeps!')
        
#     # Login to W&B
#     if relogin: wandb.login(relogin=True)
#     elif login_key: wandb.login(key=login_key)
        
#     # Sweep id
#     mod = import_file_as_module(path)
#     if not sweep_id:
#         sweep_id = wandb.sweep(mod.sweep, entity=entity, project=project)
    
#     # Agent
#     print(f"\nwandb agent {os.environ['WANDB_ENTITY']}/{os.environ['WANDB_PROJECT']}/{sweep_id}\n")
#     if launch: wandb.agent(sweep_id, function=mod.train, count=count)

# +
# export
def wandb_agent(script_path, sweep, entity=None, project=None, count=None, run=True):
    try: import wandb
    except ImportError: raise ImportError('You need to install wandb to run sweeps!')
    if 'program' not in sweep.keys(): sweep["program"] = script_path
    sweep_id = wandb.sweep(sweep, entity=entity, project=project)
    entity = ifnone(entity, os.environ['WANDB_ENTITY'])
    project = ifnone(project, os.environ['WANDB_PROJECT'])
    print(f"\nwandb agent {entity}/{project}/{sweep_id}\n")
    if run: wandb.agent(sweep_id, function=None, count=count)

get_wandb_agent = named_partial("get_wandb_agent", wandb_agent, run=False)

run_wandb_agent = named_partial("run_wandb_agent", wandb_agent, run=True)
# -

#hide
from tsai.imports import *
from tsai.export import *
nb_name = get_nb_name()
# nb_name = "201_wandb.ipynb"
create_scripts(nb_name);



StarcoderdataJupyterScriptsDedupFiltered