| |
| |
| |
| |
|
|
| import torch |
| import json |
| import numpy as np |
| from train_20m import Seq2Seq, text_to_indices |
|
|
| def load_model_and_vocab(model_path="model_20m_final.pth", vocab_path="vocab_20m.json"): |
| with open(vocab_path, "r", encoding="utf-8") as f: |
| char2idx = json.load(f) |
| idx2char = {int(v): k for k, v in char2idx.items()} |
| vocab_size = len(char2idx) |
| model = Seq2Seq(vocab_size) |
| model.load_state_dict(torch.load(model_path, map_location="cpu")) |
| model.eval() |
| return model, char2idx, idx2char |
|
|
| def generate_response(model, user_input, char2idx, idx2char, max_len=48, temperature=1.0, top_p=0.9): |
| src = text_to_indices(user_input, char2idx, max_len) |
| src_tensor = torch.tensor([src], dtype=torch.long) |
| encoder_outputs, hidden, cell = model.encoder(src_tensor) |
| decoder_input = torch.tensor([[0]], dtype=torch.long) |
| generated = [] |
| with torch.no_grad(): |
| for _ in range(64): |
| prediction, hidden, cell = model.decoder(decoder_input, encoder_outputs, hidden, cell) |
| logits = prediction[0, 0, :] / temperature |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=0), dim=0) |
| sorted_indices_to_remove = cum_probs > top_p |
| sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone() |
| sorted_indices_to_remove[0] = False |
| indices_to_remove = sorted_indices[sorted_indices_to_remove] |
| logits[indices_to_remove] = -float('Inf') |
| probs = torch.softmax(logits, dim=0).cpu().numpy() |
| next_token = np.random.choice(len(probs), p=probs) |
| if next_token == 0: |
| break |
| generated.append(idx2char[next_token]) |
| decoder_input = torch.tensor([[next_token]], dtype=torch.long) |
| return ''.join(generated) |
|
|
| if __name__ == "__main__": |
| print("加载 2000 万参数模型...") |
| model, char2idx, idx2char = load_model_and_vocab() |
| print("模型加载成功！输入 q 退出。") |
| while True: |
| user = input("你: ").strip() |
| if user.lower() == 'q': |
| break |
| reply = generate_response(model, user, char2idx, idx2char, temperature=1.1, top_p=0.9) |
| print(f"春梦蝶: {reply}") |
