topolm 0.0.4__py3-none-any.whl

topolm/__init__.py

@@ -0,0 +1,16 @@
+from .config import Config
+from .core import TopoLM, Tokenizer, Corpus, NGram, evaluate, eval_examples
+from .datasets import hf_dataset_texts, load_hf_dataset
+
+__all__ = [
+    "Config",
+    "TopoLM",
+    "Tokenizer",
+    "Corpus",
+    "NGram",
+    "evaluate",
+    "eval_examples",
+    "load_hf_dataset",
+    "hf_dataset_texts",
+]
+__version__ = "0.0.4"

topolm/cli.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import argparse
+from .core import TopoLM
+from .config import Config
+
+DEMO = """
+The cat sat on the mat.
+The dog sat on the floor.
+The attacker used CVE-2024-1234 to access the admin panel.
+CYP3A4 inhibition increases drug exposure.
+Clarithromycin inhibits CYP3A4.
+Clarithromycin may increase simvastatin exposure.
+"""
+
+def build_demo_model():
+    return TopoLM(Config()).fit(DEMO)
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="topolm")
+    sub = parser.add_subparsers(dest="cmd")
+    sub.add_parser("demo")
+    p = sub.add_parser("predict"); p.add_argument("context")
+    g = sub.add_parser("generate"); g.add_argument("prompt"); g.add_argument("--decoding", default="beam", choices=["beam", "nucleus", "greedy"])
+    args = parser.parse_args(argv)
+    model = build_demo_model()
+    if args.cmd == "predict":
+        for pred in model.distribution(args.context, 5):
+            print(f"{pred.text}\t{pred.probability:.3f}\t{pred.score:.3f}")
+    elif args.cmd == "generate":
+        print(model.generate(args.prompt, decoding=args.decoding))
+    else:
+        print(model.generate("clarithromycin inhibits", decoding="beam"))
+
+if __name__ == "__main__":
+    main()

topolm/config.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+
+PUNCT = {".", ",", ";", ":", "?", "!"}
+BOUNDARY = {"<bos>", "<eos>"}
+HUBS = {"the", "a", "an", "and", "or", "of", "to", "in", "on", "with", "when", ".", ","}
+
+@dataclass
+class Config:
+    dim: int = 1024
+    seed: int = 42
+    window: int = 8
+    phrase_lengths: tuple[int, ...] = (2, 3, 4, 5)
+    max_candidates: int = 96
+    inference_candidates: int = 48
+    prediction_cache_max: int = 4096
+    temperature: float = 0.75
+    max_runtime_seconds: float = 5.0
+    default_top_p: float = 0.88
+    default_beam_width: int = 4
+    fast_dev_mode: bool = True
+    max_reranker_sentences: int = 80
+    negatives_per_positive: int = 2
+    hub_penalty: float = 0.10