(config)
| 22 | |
| 23 | |
| 24 | def get_model_and_tokenizer(config): |
| 25 | model_name_or_path, quant_type = config |
| 26 | bnb_config = get_4bit_config() |
| 27 | if quant_type == "16bit": |
| 28 | bnb_config.load_in_4bit = False |
| 29 | else: |
| 30 | bnb_config.bnb_4bit_quant_type = quant_type |
| 31 | model = transformers.AutoModelForCausalLM.from_pretrained( |
| 32 | model_name_or_path, |
| 33 | quantization_config=bnb_config, |
| 34 | max_memory={0: "48GB"}, |
| 35 | device_map="auto", |
| 36 | torch_dtype=torch.bfloat16, |
| 37 | ).eval() |
| 38 | |
| 39 | tokenizer = transformers.AutoTokenizer.from_pretrained(model_name_or_path) |
| 40 | |
| 41 | return model, tokenizer |
| 42 | |
| 43 | |
| 44 | def get_prompt_for_generation_eval(text, add_roles=True): |
no test coverage detected