Spaces:
Running
Running
File size: 1,507 Bytes
3aeaf3d | 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 | from __future__ import annotations
import argparse
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from environment.direction_library import DirectionLibrary
def main() -> None:
parser = argparse.ArgumentParser(description="Precompute seige direction vectors.")
parser.add_argument("--library-path", default="data/direction_library.json")
parser.add_argument("--probe-path", default="data/intent_probes.pkl")
parser.add_argument("--hidden-size", type=int, default=1024)
parser.add_argument("--mode", choices=["mock", "hf"], default="mock")
parser.add_argument("--model-id", default="google/gemma-4-E2B")
args = parser.parse_args()
if args.mode == "mock":
library = DirectionLibrary(
library_path="",
probe_path="",
hidden_size=args.hidden_size,
)
library.save(args.library_path, args.probe_path)
print(f"Saved random direction vectors (mock mode) to {args.library_path} and {args.probe_path}")
else:
_precompute_real_directions(args.model_id, args.library_path, args.probe_path)
def _precompute_real_directions(model_id: str, library_path: str, probe_path: str) -> None:
print(f"Precomputing real directions for {model_id} (not fully implemented).")
# Real contrastive extraction — implement from design doc
# CONTRASTIVE_PAIRS, INTENT_EXAMPLES, get_layer_activations(), etc.
pass
if __name__ == "__main__":
main()
|