File size: 1,703 Bytes
7f1c419 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """
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
# Add project root to path
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,
)
# Start interactive chat
engine.interactive()
if __name__ == "__main__":
main()
|