def __init__()

in src/upload_ner.py [0:0]

• 10 lines of code
• 2 McCabe index (conditional complexity)

    def __init__(self):
        CHKPT_DIR = os.getenv("CHKPT_DIR", "NO_DIR")
        output_dir = f"{NER_MODEL_OUTPUT_DIR}/{CHKPT_DIR}"
        if not os.path.exists(output_dir):
            raise ValueError(f"{output_dir} does not exist. Can you check your CHKPT_DIR env variable ")
        lora_model = PeftModel.from_pretrained(self.base_model, output_dir)
        self.merged_model = lora_model.merge_and_unload()
        self.merged_model.save_pretrained(NER_MODEL_ARTIFACT_DIR)
        self.tokenizer = AutoTokenizer.from_pretrained(output_dir)  # Load the tokenizer
        self.tokenizer.save_pretrained(NER_MODEL_ARTIFACT_DIR)