| |
| |
| |
| |
|
|
| 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}") |
