File size: 2,260 Bytes
f260c78 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import json, pathlib, subprocess, sys, time
root = pathlib.Path('/content/sub1quant')
log = root / 'colab_runs' / 'setup_model_inner.log'
def emit(msg):
print(msg, flush=True)
with log.open('a', encoding='utf-8') as f:
f.write(msg + '\n')
def run(cmd):
emit('$ ' + ' '.join(cmd))
p = subprocess.run(cmd, text=True, capture_output=True)
emit('returncode=' + str(p.returncode))
if p.stdout:
emit('stdout:\n' + p.stdout[-4000:])
if p.stderr:
emit('stderr:\n' + p.stderr[-4000:])
if p.returncode != 0:
raise SystemExit(p.returncode)
emit('SETUP START ' + time.strftime('%Y-%m-%d %H:%M:%S'))
run([sys.executable, '-m', 'pip', 'install', '-q', '-U', 'huggingface_hub', 'datasets', 'safetensors', 'accelerate', 'transformers'])
data_dir = root / 'data'
data_dir.mkdir(exist_ok=True)
wiki = data_dir / 'wiki.test.txt'
if not wiki.exists() or wiki.stat().st_size < 1000:
emit('Downloading Salesforce/wikitext wikitext-2-raw-v1 test split')
from datasets import load_dataset
ds = load_dataset('Salesforce/wikitext', 'wikitext-2-raw-v1', split='test')
wiki.write_text('\n\n'.join(ds['text']), encoding='utf-8')
emit(f'wrote {wiki} bytes={wiki.stat().st_size}')
else:
emit(f'wikitext exists bytes={wiki.stat().st_size}')
model_dir = root / 'models' / 'gemma-4-E2B'
model_dir.mkdir(parents=True, exist_ok=True)
if not (model_dir / 'model.safetensors').exists():
emit('Downloading google/gemma-4-E2B-it snapshot')
from huggingface_hub import snapshot_download
snapshot_download(
repo_id='google/gemma-4-E2B-it',
local_dir=str(model_dir),
allow_patterns=['*.json', '*.model', '*.safetensors', '*.md'],
)
else:
emit(f'model exists bytes={(model_dir / "model.safetensors").stat().st_size}')
summary = {
'wiki_exists': wiki.exists(),
'wiki_size': wiki.stat().st_size if wiki.exists() else None,
'model_exists': (model_dir / 'model.safetensors').exists(),
'model_size': (model_dir / 'model.safetensors').stat().st_size if (model_dir / 'model.safetensors').exists() else None,
}
(root / 'colab_runs' / 'setup_summary.json').write_text(json.dumps(summary, indent=2), encoding='utf-8')
emit('SETUP DONE ' + json.dumps(summary))
|