Text Generation
Transformers
Safetensors
English
gpt
ww2
wwii
world-war-ii
world-war-two
second-world-war
history
military-history
historical
wikipedia
question-answering
chat
chatbot
conversational
assistant
domain-specific
small-language-model
small-model
from-scratch
trained-from-scratch
decoder-only
causal-lm
chatml
closed-book
english
rope
sentencepiece
custom_code
peacebell
| #!/usr/bin/env python3 | |
| """ | |
| Chat with this model on your own machine. | |
| pip install -r requirements.txt | |
| python run_model.py # multi-turn chat | |
| python run_model.py "When did WWII end?" # one answer, then exit | |
| python run_model.py --model <dir or repo id> # a model somewhere else | |
| By default the model is loaded from the directory this script sits in, so a | |
| plain download of the repo is enough: | |
| hf download <repo-id> --local-dir model && cd model && python run_model.py | |
| Runs on CPU (no GPU, no flash-attn needed); uses CUDA automatically when | |
| available. Greedy decoding (temperature 0) is the default and gives this | |
| model's most reliable answers. | |
| In-chat commands: | |
| /reset start a new conversation (forget the history) | |
| /temp N set the sampling temperature (0 = greedy) | |
| /length N set the maximum tokens per reply | |
| /help show the commands | |
| /quit exit (Ctrl-C or Ctrl-D also work) | |
| """ | |
| import argparse | |
| import os | |
| import sys | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| COMMANDS = """Commands: | |
| /reset start a new conversation | |
| /temp N sampling temperature (0 = greedy, the default) | |
| /length N maximum tokens per reply | |
| /help show this list | |
| /quit exit""" | |
| def resolve_model(model): | |
| """A local directory must hold a config.json; anything that looks like | |
| `user/name` is treated as a Hub repo id (downloaded on first use).""" | |
| if os.path.isdir(model): | |
| if os.path.exists(os.path.join(model, "config.json")): | |
| return model | |
| sys.exit(f"No model at {model} (no config.json there).\n" | |
| "Run this script from a downloaded model directory, or pass " | |
| "--model <directory or repo id>.") | |
| if "/" in model and not model.startswith((".", "/", "~")): | |
| return model | |
| sys.exit(f"No model at {model} (not a directory, not a repo id).") | |
| def main(): | |
| p = argparse.ArgumentParser( | |
| description="Chat with this model locally.", | |
| epilog=COMMANDS, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| p.add_argument("prompt", nargs="*", help="A single question (omit for interactive chat).") | |
| p.add_argument("--model", default=HERE, | |
| help="Model directory or Hub repo id (default: this script's directory).") | |
| p.add_argument("--temperature", type=float, default=0.0, | |
| help="0 = greedy (recommended). Higher adds errors and drift.") | |
| p.add_argument("--max-new-tokens", type=int, default=2048, | |
| help="Maximum tokens per reply (default: 2048; generation stops earlier " | |
| "when the answer ends).") | |
| p.add_argument("--fp32", action="store_true", | |
| help="Run the model in float32 on CUDA instead of the default bf16 " | |
| "(CPU is always float32).") | |
| args = p.parse_args() | |
| model_path = resolve_model(args.model) | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # bf16 on CUDA unless --fp32; CPU is always float32. | |
| dtype = torch.bfloat16 if (device == "cuda" and not args.fp32) else torch.float32 | |
| print(f"Loading {model_path} ({device}, {str(dtype).split('.')[-1]})...") | |
| tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_path, trust_remote_code=True, dtype=dtype).to(device).eval() | |
| temperature = args.temperature | |
| max_new_tokens = args.max_new_tokens | |
| def reply(messages): | |
| ids = torch.tensor(tok.build_chatml_ids(messages, add_generation_prompt=True)) | |
| out = model.chat_generate(ids, max_new_tokens=max_new_tokens, temperature=temperature) | |
| return tok.decode_response(out) | |
| if args.prompt: | |
| print(reply([{"role": "user", "content": " ".join(args.prompt)}])) | |
| return | |
| print("Ready. Type a message; /help lists the commands.\n") | |
| messages = [] # the whole conversation is re-sent each turn | |
| while True: | |
| try: | |
| text = input("You: ").strip() | |
| except (EOFError, KeyboardInterrupt): | |
| print() | |
| return | |
| if not text: | |
| continue | |
| if text.startswith("/"): | |
| cmd, _, arg = text.partition(" ") | |
| cmd = cmd.lower() | |
| if cmd == "/quit": | |
| return | |
| elif cmd == "/reset": | |
| messages = [] | |
| print("New conversation.") | |
| elif cmd == "/temp": | |
| try: | |
| temperature = float(arg) | |
| print(f"Temperature {temperature:g}" + (" (greedy)" if temperature < 0.01 else "")) | |
| except ValueError: | |
| print("Usage: /temp N (e.g. /temp 0 or /temp 0.7)") | |
| elif cmd == "/length": | |
| try: | |
| max_new_tokens = int(arg) | |
| print(f"Max tokens per reply {max_new_tokens}") | |
| except ValueError: | |
| print("Usage: /length N (e.g. /length 512)") | |
| elif cmd == "/help": | |
| print(COMMANDS) | |
| else: | |
| print(f"Unknown command {cmd}.\n{COMMANDS}") | |
| continue | |
| messages.append({"role": "user", "content": text}) | |
| try: | |
| answer = reply(messages) | |
| except KeyboardInterrupt: | |
| messages.pop() | |
| print("\n(interrupted)\n") | |
| continue | |
| messages.append({"role": "assistant", "content": answer}) | |
| print(f"\nModel: {answer}\n") | |
| if __name__ == "__main__": | |
| main() | |