Upload inference.py with huggingface_hub

inference.py

@@ -345,29 +345,13 @@ def load_packed_model(path: str) -> SSMoELMPacked:
     return model.eval()
 
 
-USER_ID      = 4
-ASSISTANT_ID = 5
-
-
-def build_chat_prompt(tok: Tokenizer, history: list[tuple[str, str]],
-                      user_input: str) -> list[int]:
-    """会話履歴 + 新しいユーザー入力 → トークン列 (instruct モデル用)"""
-    ids = [BOS_ID]
-    for user_turn, assistant_turn in history:
-        ids += [USER_ID] + tok.encode("\n" + user_turn).ids + [EOT_ID]
-        ids += [ASSISTANT_ID] + tok.encode("\n" + assistant_turn).ids + [EOT_ID]
-    ids += [USER_ID] + tok.encode("\n" + user_input).ids + [EOT_ID, ASSISTANT_ID]
-    return ids[-CTX_LEN:]
-
-
 # ── CLI ───────────────────────────────────────────────────────────────────────
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--ckpt",        default="model.safetensors", help="Path to packed safetensors")
     parser.add_argument("--tokenizer",   default="tokenizer.json",   help="Path to tokenizer.json")
-    parser.add_argument("--prompt",      default=None, help="Prompt for completion")
-    parser.add_argument("--chat",        action="store_true", help="Interactive chat mode (instruct models)")
+    parser.add_argument("--prompt",      default="The quick brown fox", help="Text prompt")
     parser.add_argument("--max-tokens",  type=int,   default=200,  help="Max new tokens to generate")
     parser.add_argument("--temperature", type=float, default=0.8,  help="Sampling temperature")
     parser.add_argument("--top-p",       type=float, default=0.9,  help="Top-p nucleus sampling")
@@ -380,40 +364,11 @@ def main():
     print(f"Memory: {packed_bytes/1024/1024:.1f} MB")
 
     tok = Tokenizer.from_file(args.tokenizer)
-
-    if args.chat:
-        # 対話モード
-        print("Chat mode — type 'exit' to quit, 'clear' to reset history")
-        print("=" * 50)
-        history: list[tuple[str, str]] = []
-        while True:
-            try:
-                user_input = input("\nYou: ").strip()
-            except (EOFError, KeyboardInterrupt):
-                break
-            if not user_input:
-                continue
-            if user_input.lower() == "exit":
-                break
-            if user_input.lower() == "clear":
-                history.clear()
-                print("--- history cleared ---")
-
-                continue
-            ids = build_chat_prompt(tok, history, user_input)
-            print("Assistant: ", end="", flush=True)
-            out = model.generate(ids, args.max_tokens, args.temperature, args.top_p)
-            response = tok.decode(out).strip()
-            print(response)
-            history.append((user_input, response))
-    else:
-        # Single-shot mode (default)
-        prompt = args.prompt or "The quick brown fox"
-        ids = [BOS_ID] + tok.encode(prompt).ids
-        print(f"\nPrompt: {prompt}")
-        print("Output: ", end="", flush=True)
-        out = model.generate(ids, args.max_tokens, args.temperature, args.top_p)
-        print(tok.decode(out))
+    ids = [BOS_ID] + tok.encode(args.prompt).ids
+    print(f"\nPrompt: {args.prompt}")
+    print("Output: ", end="", flush=True)
+    out = model.generate(ids, args.max_tokens, args.temperature, args.top_p)
+    print(tok.decode(out))
 
 
 if __name__ == "__main__":