#!/usr/bin/env python3
# Copyright (c) 2026 XingChina
# SPDX-License-Identifier: BSD-3-Clause
# 本代码采用 BSD 3-Clause 许可证，详见项目根目录的 LICENSE 文件。

import torch
import torch.nn as nn
import numpy as np
import json
import random

with open("checkpoints_lstm/char2idx.json", "r", encoding="utf-8") as f:
    char2idx = json.load(f)
idx2char = {int(v): k for k, v in char2idx.items()}
vocab_size = len(char2idx)

class CatgirlLSTM(nn.Module):
    def __init__(self, vocab_size, embed_size=128, hidden_size=256, num_layers=2, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True, dropout=dropout)
        self.fc = nn.Linear(hidden_size, vocab_size)
    def forward(self, x, hidden=None):
        x = self.embedding(x)
        out, hidden = self.lstm(x, hidden)
        out = self.fc(out)
        return out, hidden

model = CatgirlLSTM(vocab_size)
model.load_state_dict(torch.load("checkpoints_lstm/mengdie_lstm_final.pth", map_location='cpu'))
model.eval()
print("春梦蝶 LSTM 大模型加载成功！喵～\n")

def generate_response(prompt, length=2000, temperature=0.8):
    if not prompt:
        prompt = random.choice(list(char2idx.keys()))
    indices = [char2idx.get(ch, random.choice(list(char2idx.values()))) for ch in prompt]
    input_tensor = torch.tensor([indices])
    hidden = None
    result = list(prompt)
    with torch.no_grad():
        for _ in range(length):
            logits, hidden = model(input_tensor, hidden)
            probs = torch.softmax(logits[0, -1] / temperature, dim=0).cpu().numpy()
            next_idx = np.random.choice(len(probs), p=probs)
            next_char = idx2char[next_idx]
            result.append(next_char)
            input_tensor = torch.tensor([[next_idx]])
    return ''.join(result)

print("开始对话（输入 q 退出）")
while True:
    user = input("\n你: ")
    if user.lower() == 'q':
        break
    start = user[-5:] if len(user) >= 5 else user
    reply = generate_response(start, length=180, temperature=0.85)
    print(f"春梦蝶: {reply}")