I utilize the code below to run the Mistral-7b model.
However, loading the model itself is where most of the computation time is spent (Loading checkpoint shards: )

I would like to run only the inference function evaluate_model() calling python run_mycode.py, a bit like in a notebook, where you only load the model once, then you can run the model in a relatively short time. How do I do that?

from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')
import pandas as pd
import torch
import csv
from datasets import Dataset, load_dataset
from huggingface_hub import notebook_login
from peft import LoraConfig, PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
    logging,
)
from trl import SFTTrainer
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer


def evaluate_model(model, tokenizer, eval_prompt):

    model_input = tokenizer(eval_prompt, return_tensors="pt").to("cpu")

    model.eval()
    with torch.no_grad():
        return tokenizer.decode(model.generate(**model_input, max_new_tokens=256, pad_token_id=2)[0], skip_special_tokens=True)


if __name__ == "__main__":
    # base model
    torch.cuda.empty_cache()

    model_name = "mistralai/Mistral-7B-Instruct-v0.2"
    

    # Load the base model with QLoRA configuration
    compute_dtype = getattr(torch, "float16")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True, # nested quantisation (additional quantisation)
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )

    base_model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config)

    # Load MistralAi tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    
    # run model
    evaluate_model(base_model, tokenizer, eval_prompt)