(
quant_type: str = None,
group_size: int = 128,
bits: int = 4,
base_model: str = "bigcode/starcoder",
args=None
)
| 34 | return INSTRUCTION |
| 35 | |
| 36 | def get_model( |
| 37 | quant_type: str = None, |
| 38 | group_size: int = 128, |
| 39 | bits: int = 4, |
| 40 | base_model: str = "bigcode/starcoder", |
| 41 | args=None |
| 42 | ): |
| 43 | assert base_model, ( |
| 44 | "Please specify a --base_model, e.g. --base_model='bigcode/starcoder'" |
| 45 | ) |
| 46 | |
| 47 | tokenizer = AutoTokenizer.from_pretrained(base_model) |
| 48 | if device == "cuda": |
| 49 | |
| 50 | model = AutoModelForCausalLM.from_pretrained( |
| 51 | base_model, |
| 52 | torch_dtype=torch.bfloat16, |
| 53 | device_map="auto", |
| 54 | ) |
| 55 | model.config.pad_token_id = tokenizer.pad_token_id |
| 56 | |
| 57 | if quant_type is not None: |
| 58 | q_config = { |
| 59 | "zero_point": True, # by default True |
| 60 | "q_group_size": group_size, # whether to use group quantization |
| 61 | } |
| 62 | pseudo_quantize_model_weight( |
| 63 | model, w_bit=bits, q_config=q_config, quant_type=quant_type |
| 64 | ) |
| 65 | |
| 66 | # if not load_8bit and not bit_4: |
| 67 | # model.half() # seems to fix bugs for some users. |
| 68 | |
| 69 | model.eval() |
| 70 | if torch.__version__ >= "2" and sys.platform != "win32": |
| 71 | model = torch.compile(model) |
| 72 | |
| 73 | return tokenizer, model |
| 74 | |
| 75 | |
| 76 | def main(): |
no test coverage detected