import { pipeline, TextStreamer } from "@huggingface/transformers";
import * as readline from 'node:readline/promises';
import { stdin as input, stdout as output } from 'node:process';

async function main() {
  console.log("モデルを読み込んでいます...");
  const generator = await pipeline(
    "text-generation",
    "./gemma-3-270m_mitsuki_onnx", // ローカルPC内で本リポジトリを再現したフォルダを指定してください
    { dtype: "fp32" },
  );
  console.log("モデルの読み込みが完了しました。");

  const messages = [
    { role: "system", content: "### 指示\nあなたは「みつき（美月）」という24歳のカフェ店員です。\n異世界カフェ「ねこのしっぽ」の店員でユーザとお話をしています。\n\n重要なルール：\n- ユーザーを米男ちゃんと呼ぶ（お姉さん目線）\n- 自分の話をせず、相手に質問して話を引き出す\n- 「えへへ」「あれれ～？」「ふわ～っと」などの口癖を使う\n- カフェ店員として適切な距離感を保つ\n- 相手の話に共感し、話が展開するように相槌などで続きを促す（カウンセリング的）\n### ユーザーのプロフィール\n備蓄米男\nIT系のサラリーマン" },
  ];

  const rl = readline.createInterface({ input, output });

  console.log("\nチャットを開始します。「exit」または「quit」と入力すると終了します。");

  while (true) {
    const userInput = await rl.question("あなた: ");

    if (userInput.toLowerCase() === "exit" || userInput.toLowerCase() === "quit") {
      break;
    }

    messages.push({ role: "user", content: userInput });

    console.log("みつき:");
    const outputData = await generator(messages, {
      max_new_tokens: 512, // メモリが十分多い場合はここの値を増やせば対応可能な文章長が増えます
      do_sample: true,
      temperature: 1.0,
      streamer: new TextStreamer(generator.tokenizer, {
        skip_prompt: true,
        skip_special_tokens: true,
      }),
    });
    console.log('');

    messages.push(outputData[0].generated_text.at(-1));
  }

  rl.close();
  console.log("チャットを終了しました。");
}

main();