Model/scripts/chat.py

#!/usr/bin/env python3
"""
Frox AI Morph 1.1 — Interactive Chat CLI

Usage:
    python scripts/chat.py --model ./frox-morph-1-1-output/sft_final
    python scripts/chat.py --model ./frox-morph-1-1-output/sft_final --image cat.jpg
    python scripts/chat.py --model ./frox-morph-1-1-output/sft_final --quantization 4bit
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from inference.engine.morph_engine import MorphInferenceEngine
from utils.common import print_banner


SYSTEM_PROMPT = (
    "You are Frox AI Morph 1.1, a highly capable multimodal AI assistant. "
    "You can understand text, images, and video. You can search the web, "
    "run code, generate images and videos. Always be helpful, accurate, "
    "safe, and honest."
)


def main():
    parser = argparse.ArgumentParser(description="Chat with Frox AI Morph 1.1")
    parser.add_argument("--model", type=str, required=True, help="Path to saved model dir")
    parser.add_argument("--image", type=str, default=None, help="Optional image to discuss")
    parser.add_argument("--quantization", choices=["4bit", "8bit"], default=None)
    parser.add_argument("--max-tokens", type=int, default=1024)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--session", action="store_true",
                        help="Use persistent KV-cache session mode (faster multi-turn)")
    args = parser.parse_args()

    print_banner()
    engine = MorphInferenceEngine.from_pretrained(
        args.model, quantization=args.quantization,
    )

    pixel_values = None
    if args.image:
        from PIL import Image
        img = Image.open(args.image).convert("RGB")
        pixel_values = engine.model.vision_module.preprocess_image(img, device=engine.device)
        print(f"🖼  Loaded image: {args.image}\n")

    print("Type your message. Commands: /reset  /think  /quit\n")

    messages = []
    session_id = "cli-session"
    show_thinking = False

    while True:
        try:
            user_input = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if not user_input:
            continue
        if user_input == "/quit":
            break
        if user_input == "/reset":
            messages = []
            engine.reset_session(session_id)
            print("(conversation reset)\n")
            continue
        if user_input == "/think":
            show_thinking = not show_thinking
            print(f"(thinking display: {'on' if show_thinking else 'off'})\n")
            continue

        print("Morph: ", end="", flush=True)

        if args.session:
            response = engine.chat(
                session_id, user_input, system_prompt=SYSTEM_PROMPT,
                max_new_tokens=args.max_tokens, temperature=args.temperature,
                stream_callback=lambda chunk: print(chunk, end="", flush=True),
            )
            print()
        else:
            messages.append({"role": "user", "content": user_input})
            response_parts = []
            for chunk in engine.generate_stream(
                messages, system_prompt=SYSTEM_PROMPT,
                max_new_tokens=args.max_tokens, temperature=args.temperature,
                pixel_values=pixel_values,
            ):
                print(chunk, end="", flush=True)
                response_parts.append(chunk)
            response = "".join(response_parts)
            print()
            messages.append({"role": "assistant", "content": response})

        thinking, answer = engine.parse_thinking(response)
        if thinking and show_thinking:
            print(f"\n💭 [thinking]: {thinking}\n")

        tool_calls = engine.parse_tool_calls(response)
        if tool_calls:
            print(f"\n🔧 [tool calls detected]: {tool_calls}")

        gen_requests = engine.parse_generation_requests(response)
        if gen_requests:
            print(f"\n🎨 [generation requests]: {gen_requests}")

        print()


if __name__ == "__main__":
    main()