def __init__()

in evolve-instruct/evolve.py [0:0]

• 21 lines of code
• 1 McCabe index (conditional complexity)

    def __init__(self, model, max_new_tokens=None, batch_size=None, **kwargs):
        from transformers import AutoTokenizer, AutoModelForCausalLM
        print("loading tokenizer")
        tokenizer = AutoTokenizer.from_pretrained(model, padding_side="left")
        print("loading model")
        model_obj = AutoModelForCausalLM.from_pretrained(model, torch_dtype=torch.bfloat16, device_map="auto")
        pad_token_id = model_obj.config.eos_token_id
        del model_obj
        print("loading pipeline")
        self.pipeline = pipeline(
            model=model,
            tokenizer=tokenizer,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            device_map="auto",
            **kwargs,
        )
        print("loading pipeline done.")
        self.pipeline.tokenizer.pad_token_id = pad_token_id
        self.max_new_tokens = max_new_tokens
        self.batch_size = batch_size