refactor(prefetch): clean up unused variables, dead functions, and hardcoded paths
This commit is contained in:
@@ -22,29 +22,25 @@ try:
|
||||
except ImportError as exc:
|
||||
sys.exit(f"Missing deps: {exc}. Run: pip install torch transformers")
|
||||
|
||||
MODEL_DIR = (
|
||||
Path(r"C:\Users\egonr\.cache\huggingface\hub"
|
||||
r"\models--allenai--OLMoE-1B-7B-0125-Instruct"
|
||||
r"\snapshots\b89a7c4bc24fb9e55ce2543c9458ce0ca5c4650e")
|
||||
)
|
||||
MODEL_ID = "allenai/OLMoE-1B-7B-0125-Instruct"
|
||||
|
||||
OUT_JSON = Path(__file__).resolve().parent.parent / "ref_olmoe_real.json"
|
||||
|
||||
PROMPT = "The capital of France is"
|
||||
MAX_NEW_TOKENS = 12
|
||||
|
||||
print(f"Loading tokenizer from {MODEL_DIR} ...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR))
|
||||
print(f"Loading tokenizer from {MODEL_ID} ...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
||||
|
||||
print("Encoding prompt ...")
|
||||
enc = tokenizer(PROMPT, return_tensors="pt")
|
||||
prompt_ids = enc["input_ids"][0].tolist()
|
||||
print(f" Prompt IDs ({len(prompt_ids)}): {prompt_ids}")
|
||||
|
||||
print(f"Loading OLMoE model from {MODEL_DIR} ...")
|
||||
print(f"Loading OLMoE model from {MODEL_ID} ...")
|
||||
print(" (this will use ~14 GB RAM — please be patient)")
|
||||
model = OlmoeForCausalLM.from_pretrained(
|
||||
str(MODEL_DIR),
|
||||
MODEL_ID,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="cpu",
|
||||
low_cpu_mem_usage=True,
|
||||
|
||||
@@ -22,8 +22,9 @@ if sys.platform == "win32":
|
||||
pass
|
||||
|
||||
HERE = Path(__file__).resolve().parent.parent
|
||||
ENGINE = HERE / "olmoe.exe"
|
||||
SNAP = HERE / "olmoe_i4"
|
||||
ext = ".exe" if sys.platform == "win32" else ""
|
||||
ENGINE = HERE / f"olmoe{ext}"
|
||||
SNAP = os.getenv("SNAP", str(HERE.parent / "olmoe_merged"))
|
||||
REF_OUT = HERE / "ref_olmoe_real.json"
|
||||
BOOTSTRAP_REF = HERE / "ref_olmoe_bootstrap.json"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user