| """ |
| ViuAI-1M Chat Script. |
| |
| Start an interactive chat session with a trained ViuAI model. |
| |
| Usage: |
| python scripts/chat.py --checkpoint output/final --tokenizer output/tokenizer |
| """ |
|
|
| import argparse |
| import os |
| import sys |
|
|
| |
| project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, project_root) |
|
|
| from viuai.inference.chat import load_model_for_chat |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description="Chat with ViuAI-1M") |
| parser.add_argument("--checkpoint", type=str, default="output/final", |
| help="Path to model checkpoint directory") |
| parser.add_argument("--tokenizer", type=str, default="output/tokenizer", |
| help="Path to tokenizer directory") |
| parser.add_argument("--device", type=str, default=None, |
| help="Device (cuda/cpu)") |
| parser.add_argument("--temperature", type=float, default=0.8, |
| help="Sampling temperature") |
| parser.add_argument("--max_tokens", type=int, default=128, |
| help="Maximum tokens to generate") |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
|
|
| checkpoint_path = os.path.join(project_root, args.checkpoint) |
| tokenizer_path = os.path.join(project_root, args.tokenizer) |
|
|
| print(f"\n Loading ViuAI-1M...") |
| print(f" Checkpoint: {checkpoint_path}") |
| print(f" Tokenizer: {tokenizer_path}") |
|
|
| engine = load_model_for_chat( |
| checkpoint_dir=checkpoint_path, |
| tokenizer_dir=tokenizer_path, |
| device=args.device, |
| ) |
|
|
| |
| engine.interactive() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
