neuralbroker commited on
Commit
0a72c46
·
verified ·
1 Parent(s): cd53253

Update clean backend-only project docs and eval

Browse files
Files changed (1) hide show
  1. scripts/train_available.py +3 -3
scripts/train_available.py CHANGED
@@ -12,7 +12,7 @@ import argparse
12
  import json
13
  import os
14
  from pathlib import Path
15
- from typing import Any
16
 
17
  import torch
18
  from datasets import Dataset
@@ -194,14 +194,14 @@ def load_model(model_name: str, quantization: str) -> tuple[Any, Any, bool]:
194
  quantization_config = BitsAndBytesConfig(
195
  load_in_4bit=True,
196
  bnb_4bit_quant_type="nf4",
197
- bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
198
  bnb_4bit_use_double_quant=True,
199
  )
200
  except Exception as exc:
201
  print(f"[WARN] 4-bit config unavailable, falling back to 16-bit load: {exc}")
202
  use_4bit = False
203
 
204
- dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
205
  model_kwargs: dict[str, Any] = {
206
  "trust_remote_code": True,
207
  "device_map": "auto" if torch.cuda.is_available() else None,
 
12
  import json
13
  import os
14
  from pathlib import Path
15
+ from typing import Any, cast
16
 
17
  import torch
18
  from datasets import Dataset
 
194
  quantization_config = BitsAndBytesConfig(
195
  load_in_4bit=True,
196
  bnb_4bit_quant_type="nf4",
197
+ bnb_4bit_compute_dtype=cast(Any, torch).bfloat16 if torch.cuda.is_bf16_supported() else cast(Any, torch).float16,
198
  bnb_4bit_use_double_quant=True,
199
  )
200
  except Exception as exc:
201
  print(f"[WARN] 4-bit config unavailable, falling back to 16-bit load: {exc}")
202
  use_4bit = False
203
 
204
+ dtype = cast(Any, torch).bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else cast(Any, torch).float16
205
  model_kwargs: dict[str, Any] = {
206
  "trust_remote_code": True,
207
  "device_map": "auto" if torch.cuda.is_available() else None,