| 9 | vocab_filename = "bpe_simple_vocab_16e6.txt" |
| 10 | |
| 11 | def ensure_data_dir(safetensors): |
| 12 | print("Ensuring empty data directory...") |
| 13 | |
| 14 | if os.path.exists(data_path): |
| 15 | # Fail if conflicting files exist |
| 16 | files = os.listdir(data_path) |
| 17 | newfiles = [x for name in safetensors for x in (f"{name}.bin", f"{name}.safetensors")] |
| 18 | newfiles += [vocab_filename, f"{vocab_filename}.gz"] |
| 19 | conflicts = set(files) & set(newfiles) |
| 20 | if len(conflicts) != 0: |
| 21 | print("Error: please remove the following files from data directory:") |
| 22 | print(conflicts) |
| 23 | sys.exit("Found conflicting files in data directory.") |
| 24 | else: |
| 25 | os.mkdir(data_path) |
| 26 | |
| 27 | print("Found no conflicts!") |
| 28 | |
| 29 | def get_safetensors(safetensors, weight_bits): |
| 30 | for name, url in safetensors.items(): |