Upload generate.py with huggingface_hub

generate.py

@@ -38,6 +38,8 @@ def main():
     parser.add_argument("--seed", type=int, default=None)
     parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
     parser.add_argument("--guard", default=None, help="file kamus (satu kata per baris); pilih kandidat dengan ejaan terbaik")
+    parser.add_argument("--guard-prefiks", default=None, help="daftar prefiks (default data/prefiks.txt bila ada)")
+    parser.add_argument("--guard-sufiks", default=None, help="daftar sufiks (default data/sufiks.txt bila ada)")
     parser.add_argument("--guard-tries", type=int, default=5, help="jumlah kandidat saat --guard aktif")
     parser.add_argument("--guard-min", type=float, default=0.6, help="rasio kata dikenal minimum (info saja)")
     args = parser.parse_args()
@@ -50,11 +52,24 @@ def main():
     model = model.to(device)
 
     wordset = None
+    pref_set = suf_set = None
     if args.guard:
+        from pathlib import Path as _Path
+
         from indigo.common import load_wordlist, word_known_ratio
 
         wordset = load_wordlist(args.guard)
-        print(f"[guard] kamus: {len(wordset):,} kata | target rasio >= {args.guard_min:.0%}")
+        p_def, s_def = _Path("data/prefiks.txt"), _Path("data/sufiks.txt")
+        if args.guard_prefiks and _Path(args.guard_prefiks).exists():
+            pref_set = load_wordlist(args.guard_prefiks)
+        elif not args.guard_prefiks and p_def.exists():
+            pref_set = load_wordlist(str(p_def))
+        if args.guard_sufiks and _Path(args.guard_sufiks).exists():
+            suf_set = load_wordlist(args.guard_sufiks)
+        elif not args.guard_sufiks and s_def.exists():
+            suf_set = load_wordlist(str(s_def))
+        mode = "dengan formula afiks" if pref_set and suf_set else "kata persis"
+        print(f"[guard] kamus: {len(wordset):,} kata ({mode}) | target rasio >= {args.guard_min:.0%}")
 
     ids = tokenizer.encode(args.prompt) or [0]
     idx = torch.tensor([ids], dtype=torch.long, device=device)
@@ -69,7 +84,7 @@ def main():
             repetition_penalty=args.repetition_penalty,
         )
         text = tokenizer.decode(out[0].tolist())
-        ratio = word_known_ratio(text, wordset) if wordset else 1.0
+        ratio = word_known_ratio(text, wordset, pref_set, suf_set) if wordset else 1.0
         return text, ratio
 
     if wordset is None: