import nltk
from nltk.tokenize import word_tokenize, sent_tokenize
nltk.download('punkt')

text = "Natural Language Processing is fascinating! It powers ChatGPT and many AI applications."

# 单词分词
word_tokens = word_tokenize(text)
print("单词词元:", word_tokens)
# 输出: ['Natural', 'Language', 'Processing', 'is', 'fascinating', '!', ...]

# 句子分词
sent_tokens = sent_tokenize(text)
print("句子词元:", sent_tokens)
# 输出: ['Natural Language Processing is fascinating!', 'It powers ChatGPT...']

text = "Hello NLP"
char_tokens = list(text)
print("字符词元:", char_tokens)
# 输出: ['H', 'e', 'l', 'l', 'o', ' ', 'N', 'L', 'P']

from tokenizers import ByteLevelBPETokenizer

# 训练 BPE 分词器
tokenizer = ByteLevelBPETokenizer()
tokenizer.train(
    files=["corpus.txt"],
    vocab_size=5000,
    min_frequency=2,
    special_tokens=["<s>", "<pad>", "</s>", "<unk>", "<mask>"]
)

# 编码
encoding = tokenizer.encode("Natural Language Processing")
print("词元:", encoding.tokens)
print("ID:", encoding.ids)

import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
nltk.download('stopwords')

text = "This is a sample sentence showing off the stop words filtration."
stop_words = set(stopwords.words('english'))

word_tokens = word_tokenize(text)
filtered_sentence = [w for w in word_tokens if not w.lower() in stop_words]

print("原始:", word_tokens)
print("过滤后:", filtered_sentence)
# 输出: ['sample', 'sentence', 'showing', 'stop', 'words', 'filtration', '.']

from nltk.stem import PorterStemmer, LancasterStemmer

ps = PorterStemmer()
ls = LancasterStemmer()

words = ["running", "runs", "ran", "runner", "easily", "fairly"]
for word in words:
    print(f"{word:15} -> Porter: {ps.stem(word):15} Lancaster: {ls.stem(word)}")

import nltk
from nltk.stem import WordNetLemmatizer
nltk.download('wordnet')

lemmatizer = WordNetLemmatizer()

# 指定词性可以得到更准确的结果
print(lemmatizer.lemmatize("running", pos='v'))  # run
print(lemmatizer.lemmatize("better", pos='a'))    # good
print(lemmatizer.lemmatize("dogs"))               # dog
print(lemmatizer.lemmatize("went", pos='v'))      # go

import re

def clean_text(text):
    """文本清洗函数"""
    # 去除 HTML 标签
    text = re.sub(r'<[^>]+>', '', text)

    # 去除 URL
    text = re.sub(r'http\S+|www\S+|https\S+', '', text, flags=re.MULTILINE)

    # 去除邮箱
    text = re.sub(r'\S+@\S+', '', text)

    # 去除特殊字符(只保留字母、数字、空格)
    text = re.sub(r'[^a-zA-Z0-9\s]', '', text)

    # 多个空格合并为一个
    text = re.sub(r'\s+', ' ', text)

    # 去除首尾空格
    text = text.strip().lower()

    return text

sample_text = """
Check out my website at https://example.com!
Email me at user@example.com for <b>more info</b>.
It's   really   cool!!!
"""

cleaned = clean_text(sample_text)
print(cleaned)
# 输出: "check out my website at email me at for more info its really cool"

import spacy

# 加载英语模型
nlp = spacy.load("en_core_web_sm")

text = "Apple is looking at buying U.K. startup for $1 billion"
doc = nlp(text)

print("=== 词元信息 ===")
for token in doc:
    print(f"{token.text:15} | POS: {token.pos_:10} | Lemma: {token.lemma_:15} | Stop: {token.is_stop}")

print("\n=== 命名实体识别 ===")
for ent in doc.ents:
    print(f"{ent.text:20} -> {ent.label_} ({spacy.explain(ent.label_)})")

print("\n=== 依存句法分析 ===")
for token in doc:
    print(f"{token.text:15} -> {token.dep_:15} (head: {token.head.text})")

from konlpy.tag import Okt, Mecab, Komoran, Kkma

# Okt (Open Korean Text) 使用示例
okt = Okt()

text = "자연어 처리는 인공지능의 핵심 분야입니다."

print("形态素分析:", okt.morphs(text))
print("词性标注:", okt.pos(text))
print("名词提取:", okt.nouns(text))
print("词组切分:", okt.phrases(text))

# 输出示例:
# 形态素分析: ['자연어', '처리', '는', '인공지능', '의', '핵심', '분야', '입니다', '.']
# 词性标注: [('자연어', 'Noun'), ('처리', 'Noun'), ('는', 'Josa'), ...]
# 名词提取: ['자연어', '처리', '인공지능', '핵심', '분야']

# Mecab —— 速度最快、最精准的形态素分析器
try:
    mecab = Mecab()
    print("\nMecab 分析:")
    print(mecab.pos(text))
except Exception:
    print("未安装 Mecab。")

# 韩语停用词处理
korean_stopwords = ['은', '는', '이', '가', '을', '를', '의', '에', '에서', '으로', '와', '과', '도']

def korean_tokenize(text):
    okt = Okt()
    tokens = okt.morphs(text, stem=True)  # 包含词干提取
    tokens = [t for t in tokens if t not in korean_stopwords]
    tokens = [t for t in tokens if len(t) > 1]  # 去除长度为1的词
    return tokens

print("\n韩语分词结果:", korean_tokenize(text))

from sklearn.feature_extraction.text import CountVectorizer
import numpy as np

corpus = [
    "I love natural language processing",
    "Natural language processing is amazing",
    "I love machine learning too",
    "Deep learning is part of machine learning"
]

# BoW 向量化器
vectorizer = CountVectorizer()
X = vectorizer.fit_transform(corpus)

print("词表:", vectorizer.get_feature_names_out())
print("\nBoW 矩阵:")
print(X.toarray())
print("\n形状:", X.shape)  # (4 篇文档, n 个单词)

# 查看指定文档
doc_0 = dict(zip(vectorizer.get_feature_names_out(), X.toarray()[0]))
print("\n文档 0 的单词频率:")
for word, count in sorted(doc_0.items()):
    if count > 0:
        print(f"  {word}: {count}")

from sklearn.feature_extraction.text import TfidfVectorizer
import pandas as pd

corpus = [
    "the cat sat on the mat",
    "the cat sat on the hat",
    "the dog sat on the log",
    "the cat wore the hat",
]

# TF-IDF 向量化器
tfidf = TfidfVectorizer(smooth_idf=True, norm='l2')
X = tfidf.fit_transform(corpus)

# 把结果放进 DataFrame 展示
df = pd.DataFrame(
    X.toarray(),
    columns=tfidf.get_feature_names_out(),
    index=[f"Doc {i}" for i in range(len(corpus))]
)
print(df.round(3))

# TF-IDF 手动计算示例
def compute_tfidf(corpus):
    from math import log

    # 构建词表
    vocab = set()
    for doc in corpus:
        vocab.update(doc.split())
    vocab = sorted(vocab)

    # 计算 TF
    def tf(word, doc):
        words = doc.split()
        return words.count(word) / len(words)

    # 计算 IDF
    def idf(word, corpus):
        n_docs_with_word = sum(1 for doc in corpus if word in doc.split())
        return log((1 + len(corpus)) / (1 + n_docs_with_word)) + 1

    # 计算 TF-IDF 矩阵
    tfidf_matrix = []
    for doc in corpus:
        row = [tf(word, doc) * idf(word, corpus) for word in vocab]
        tfidf_matrix.append(row)

    return tfidf_matrix, vocab

matrix, vocab = compute_tfidf(corpus)
print("\n手动计算的 TF-IDF:")
df_manual = pd.DataFrame(matrix, columns=vocab)
print(df_manual.round(3))

from nltk.util import ngrams
from nltk.tokenize import word_tokenize
from collections import Counter

text = "I love natural language processing and machine learning"
tokens = word_tokenize(text.lower())

# Unigram (1-gram)
unigrams = list(ngrams(tokens, 1))
# Bigram (2-gram)
bigrams = list(ngrams(tokens, 2))
# Trigram (3-gram)
trigrams = list(ngrams(tokens, 3))

print("Unigrams:", unigrams[:5])
print("Bigrams:", bigrams[:5])
print("Trigrams:", trigrams[:5])

# 计算 N-gram 频率
def get_ngram_freq(text, n):
    tokens = word_tokenize(text.lower())
    n_grams = list(ngrams(tokens, n))
    return Counter(n_grams)

# 在更大语料上的应用
large_corpus = """
Machine learning is a subset of artificial intelligence.
Artificial intelligence is transforming many industries.
Natural language processing is a part of machine learning.
Deep learning has revolutionized natural language processing.
"""

bigram_freq = get_ngram_freq(large_corpus, 2)
print("\n出现频率最高的 Bigrams:")
for bigram, count in bigram_freq.most_common(10):
    print(f"  {' '.join(bigram)}: {count}")

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import Counter

class Word2VecSkipGram(nn.Module):
    """Skip-gram Word2Vec 实现"""

    def __init__(self, vocab_size, embedding_dim):
        super().__init__()
        self.center_embedding = nn.Embedding(vocab_size, embedding_dim)
        self.context_embedding = nn.Embedding(vocab_size, embedding_dim)

        # 初始化
        self.center_embedding.weight.data.uniform_(-0.5/embedding_dim, 0.5/embedding_dim)
        self.context_embedding.weight.data.uniform_(-0.5/embedding_dim, 0.5/embedding_dim)

    def forward(self, center, context, negative):
        """
        center: (batch_size,) 中心词索引
        context: (batch_size,) 实际上下文词索引
        negative: (batch_size, neg_samples) 负采样索引
        """
        # 中心词嵌入
        center_emb = self.center_embedding(center)  # (batch, dim)

        # 正样本得分
        context_emb = self.context_embedding(context)  # (batch, dim)
        pos_score = torch.sum(center_emb * context_emb, dim=1)  # (batch,)
        pos_loss = -torch.log(torch.sigmoid(pos_score) + 1e-10)

        # 负样本得分
        neg_emb = self.context_embedding(negative)  # (batch, neg_samples, dim)
        center_emb_expanded = center_emb.unsqueeze(1)  # (batch, 1, dim)
        neg_score = torch.bmm(neg_emb, center_emb_expanded.transpose(1, 2)).squeeze(2)  # (batch, neg_samples)
        neg_loss = -torch.sum(torch.log(torch.sigmoid(-neg_score) + 1e-10), dim=1)

        return (pos_loss + neg_loss).mean()

class Word2VecCBOW(nn.Module):
    """CBOW Word2Vec 实现"""

    def __init__(self, vocab_size, embedding_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.linear = nn.Linear(embedding_dim, vocab_size)

    def forward(self, context):
        """
        context: (batch_size, context_window_size) 上下文词索引
        """
        # 上下文词嵌入的平均值
        embedded = self.embedding(context)  # (batch, window, dim)
        mean_embedded = embedded.mean(dim=1)  # (batch, dim)
        output = self.linear(mean_embedded)  # (batch, vocab_size)
        return output

# 数据准备与训练示例
def prepare_data(sentences, window_size=2):
    """准备训练数据"""
    word_counts = Counter()
    for sentence in sentences:
        word_counts.update(sentence.split())

    vocab = ['<UNK>'] + [word for word, _ in word_counts.most_common()]
    word2idx = {word: idx for idx, word in enumerate(vocab)}
    idx2word = {idx: word for word, idx in word2idx.items()}

    training_data = []
    for sentence in sentences:
        words = sentence.split()
        for i, center_word in enumerate(words):
            center_idx = word2idx.get(center_word, 0)
            for j in range(max(0, i-window_size), min(len(words), i+window_size+1)):
                if i != j:
                    context_idx = word2idx.get(words[j], 0)
                    training_data.append((center_idx, context_idx))

    return training_data, word2idx, idx2word, vocab

# 使用 Gensim 实现 Word2Vec(更实用的方法)
from gensim.models import Word2Vec

sentences = [
    "I love natural language processing".split(),
    "natural language processing is a field of AI".split(),
    "machine learning is part of artificial intelligence".split(),
    "deep learning models process natural language".split(),
    "word embeddings represent words as vectors".split(),
    "Word2Vec learns word representations".split(),
    "semantic similarity between words is captured by embeddings".split(),
]

# 训练模型
model = Word2Vec(
    sentences=sentences,
    vector_size=100,   # 嵌入维度
    window=5,          # 窗口大小
    min_count=1,       # 最小出现次数
    workers=4,         # 并行处理数
    epochs=100,        # 训练轮数
    sg=1,              # 1=Skip-gram, 0=CBOW
    negative=5         # 负采样数
)

# 查看单词向量
print("'language' 向量(前 5 维):", model.wv['language'][:5])

# 查找相似单词
similar_words = model.wv.most_similar('language', topn=5)
print("\n与 'language' 相似的单词:")
for word, similarity in similar_words:
    print(f"  {word}: {similarity:.4f}")

# 单词类比(King - Man + Woman = Queen 类型)
# 类比: language - natural + artificial = ?
result = model.wv.most_similar(
    positive=['artificial', 'language'],
    negative=['natural'],
    topn=3
)
print("\n单词类比结果:")
for word, sim in result:
    print(f"  {word}: {sim:.4f}")

# 两个单词的相似度
print("\n'language' 与 'processing' 的相似度:", model.wv.similarity('language', 'processing'))

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np

def visualize_embeddings(model, words=None):
    """Word2Vec 嵌入的 t-SNE 可视化"""
    if words is None:
        words = list(model.wv.key_to_index.keys())[:50]

    # 提取向量
    vectors = np.array([model.wv[word] for word in words])

    # t-SNE 降维
    tsne = TSNE(
        n_components=2,
        random_state=42,
        perplexity=min(30, len(words)-1),
        n_iter=1000
    )
    vectors_2d = tsne.fit_transform(vectors)

    # 可视化
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.scatter(vectors_2d[:, 0], vectors_2d[:, 1], alpha=0.7)

    for i, word in enumerate(words):
        ax.annotate(word, (vectors_2d[i, 0], vectors_2d[i, 1]),
                   fontsize=9, ha='center', va='bottom')

    ax.set_title("Word2Vec Embeddings t-SNE Visualization")
    plt.tight_layout()
    plt.savefig('word2vec_tsne.png', dpi=150, bbox_inches='tight')
    plt.show()

visualize_embeddings(model)

# 使用预训练的 GloVe 向量
# pip install torchtext

from torchtext.vocab import GloVe

# 下载并加载 GloVe
glove = GloVe(name='6B', dim=100)

# 访问单词向量
vector = glove['computer']
print("'computer' 的 GloVe 向量(前 5 维):", vector[:5].numpy())

# 计算余弦相似度
def cosine_similarity(v1, v2):
    return torch.nn.functional.cosine_similarity(
        v1.unsqueeze(0), v2.unsqueeze(0)
    ).item()

import torch
words_to_compare = [('king', 'queen'), ('man', 'woman'), ('Paris', 'France')]
for w1, w2 in words_to_compare:
    if w1 in glove.stoi and w2 in glove.stoi:
        sim = cosine_similarity(glove[w1], glove[w2])
        print(f"sim('{w1}', '{w2}') = {sim:.4f}")

import torch
import torch.nn as nn

class SimpleRNN(nn.Module):
    """基本 RNN 实现"""

    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.hidden_size = hidden_size

        # 输入 -> 隐藏状态 转换
        self.input_to_hidden = nn.Linear(input_size + hidden_size, hidden_size)
        # 隐藏状态 -> 输出 转换
        self.hidden_to_output = nn.Linear(hidden_size, output_size)

    def forward(self, x, hidden=None):
        """
        x: (seq_len, batch, input_size)
        hidden: (batch, hidden_size) 初始隐藏状态
        """
        batch_size = x.size(1)

        if hidden is None:
            hidden = torch.zeros(batch_size, self.hidden_size)

        outputs = []
        for t in range(x.size(0)):
            # 拼接输入与前一个隐藏状态
            combined = torch.cat([x[t], hidden], dim=1)
            # 计算新的隐藏状态
            hidden = torch.tanh(self.input_to_hidden(combined))
            # 计算输出
            output = self.hidden_to_output(hidden)
            outputs.append(output)

        return torch.stack(outputs, dim=0), hidden

# 使用 PyTorch 内置 RNN
rnn = nn.RNN(
    input_size=50,
    hidden_size=128,
    num_layers=2,
    batch_first=True,
    dropout=0.3,
    bidirectional=True  # 双向 RNN
)

# 用随机输入测试
x = torch.randn(32, 20, 50)  # (batch, seq_len, features)
output, hidden = rnn(x)
print("输出形状:", output.shape)   # (32, 20, 256) - 双向所以是 128*2
print("隐藏状态形状:", hidden.shape)  # (4, 32, 128) - 2层 * 2方向

class LSTMCell(nn.Module):
    """手动实现 LSTM 单元 —— 用于理解门控机制"""

    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size

        # 一次性计算 4 个门(提升效率)
        # forget gate, input gate, cell gate, output gate
        self.gates = nn.Linear(input_size + hidden_size, 4 * hidden_size)

    def forward(self, x, state=None):
        """
        x: (batch, input_size)
        state: (h, c) 前一个隐藏状态与细胞状态
        """
        batch_size = x.size(0)

        if state is None:
            h = torch.zeros(batch_size, self.hidden_size)
            c = torch.zeros(batch_size, self.hidden_size)
        else:
            h, c = state

        # 拼接输入与前一个隐藏状态
        combined = torch.cat([x, h], dim=1)

        # 计算 4 个门
        gates = self.gates(combined)

        # 拆分门
        f_gate, i_gate, g_gate, o_gate = gates.chunk(4, dim=1)

        # 对每个门应用激活函数
        f = torch.sigmoid(f_gate)   # forget gate: 决定遗忘多少旧信息
        i = torch.sigmoid(i_gate)   # input gate: 决定存储多少新信息
        g = torch.tanh(g_gate)      # cell gate: 新的候选信息
        o = torch.sigmoid(o_gate)   # output gate: 决定输出多少

        # 新的细胞状态: 部分遗忘 + 加入新信息
        new_c = f * c + i * g

        # 新的隐藏状态
        new_h = o * torch.tanh(new_c)

        return new_h, (new_h, new_c)

# 使用 PyTorch 内置 LSTM 的情感分析模型
class SentimentLSTM(nn.Module):
    """基于 LSTM 的情感分析模型"""

    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes, num_layers=2, dropout=0.3):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            embed_dim,
            hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout,
            bidirectional=True
        )
        self.dropout = nn.Dropout(dropout)
        # 双向所以是 hidden_dim * 2
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes)
        )

    def forward(self, x, lengths=None):
        """
        x: (batch, seq_len) 词元索引
        lengths: 每个序列的实际长度(考虑填充)
        """
        # 嵌入
        embedded = self.dropout(self.embedding(x))  # (batch, seq, embed)

        # 处理填充序列
        if lengths is not None:
            packed = nn.utils.rnn.pack_padded_sequence(
                embedded, lengths, batch_first=True, enforce_sorted=False
            )
            lstm_out, (hidden, cell) = self.lstm(packed)
            lstm_out, _ = nn.utils.rnn.pad_packed_sequence(lstm_out, batch_first=True)
        else:
            lstm_out, (hidden, cell) = self.lstm(embedded)

        # 最后一层的隐藏状态(双向)
        # hidden: (num_layers * 2, batch, hidden_dim)
        forward_hidden = hidden[-2]   # 最后一层的正向
        backward_hidden = hidden[-1]  # 最后一层的反向
        combined = torch.cat([forward_hidden, backward_hidden], dim=1)

        # 分类
        output = self.classifier(self.dropout(combined))
        return output

class GRUCell(nn.Module):
    """手动实现 GRU 单元"""

    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size

        # Reset gate
        self.reset_gate = nn.Linear(input_size + hidden_size, hidden_size)
        # Update gate
        self.update_gate = nn.Linear(input_size + hidden_size, hidden_size)
        # New gate
        self.new_gate_input = nn.Linear(input_size, hidden_size)
        self.new_gate_hidden = nn.Linear(hidden_size, hidden_size)

    def forward(self, x, h=None):
        batch_size = x.size(0)

        if h is None:
            h = torch.zeros(batch_size, self.hidden_size)

        combined = torch.cat([x, h], dim=1)

        # 重置门: 决定重置多少旧隐藏状态
        r = torch.sigmoid(self.reset_gate(combined))

        # 更新门: 决定更新多少
        z = torch.sigmoid(self.update_gate(combined))

        # 新的候选状态
        n = torch.tanh(
            self.new_gate_input(x) + r * self.new_gate_hidden(h)
        )

        # 新的隐藏状态
        new_h = (1 - z) * n + z * h

        return new_h

class Encoder(nn.Module):
    """Seq2Seq 编码器"""

    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers=2, dropout=0.5):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers,
                           batch_first=True, dropout=dropout)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        embedded = self.dropout(self.embedding(src))
        outputs, (hidden, cell) = self.lstm(embedded)
        return outputs, hidden, cell

class Decoder(nn.Module):
    """Seq2Seq 解码器"""

    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers=2, dropout=0.5):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers,
                           batch_first=True, dropout=dropout)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, trg, hidden, cell):
        # trg: (batch,) 当前词元
        trg = trg.unsqueeze(1)  # (batch, 1)
        embedded = self.dropout(self.embedding(trg))  # (batch, 1, embed)
        output, (hidden, cell) = self.lstm(embedded, (hidden, cell))
        prediction = self.fc(output.squeeze(1))  # (batch, vocab_size)
        return prediction, hidden, cell

class Seq2Seq(nn.Module):
    """Seq2Seq 翻译模型"""

    def __init__(self, encoder, decoder, device):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.device = device

    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        """
        src: (batch, src_len) 源序列
        trg: (batch, trg_len) 目标序列
        teacher_forcing_ratio: teacher forcing 概率
        """
        batch_size = src.shape[0]
        trg_len = trg.shape[1]
        trg_vocab_size = self.decoder.fc.out_features

        outputs = torch.zeros(batch_size, trg_len, trg_vocab_size).to(self.device)

        # 运行编码器
        encoder_outputs, hidden, cell = self.encoder(src)

        # 解码器的第一个输入: <SOS> 词元
        input = trg[:, 0]

        for t in range(1, trg_len):
            output, hidden, cell = self.decoder(input, hidden, cell)
            outputs[:, t] = output

            # teacher forcing: 训练时使用真实标签或预测值
            teacher_force = torch.rand(1).item() < teacher_forcing_ratio
            top1 = output.argmax(1)
            input = trg[:, t] if teacher_force else top1

        return outputs

class BahdanauAttention(nn.Module):
    """Bahdanau (Additive) Attention"""

    def __init__(self, hidden_dim):
        super().__init__()
        self.W_enc = nn.Linear(hidden_dim * 2, hidden_dim)  # 双向编码器
        self.W_dec = nn.Linear(hidden_dim, hidden_dim)
        self.v = nn.Linear(hidden_dim, 1)

    def forward(self, decoder_hidden, encoder_outputs):
        """
        decoder_hidden: (batch, hidden) 解码器当前的隐藏状态
        encoder_outputs: (batch, src_len, hidden*2) 全部编码器输出
        """
        src_len = encoder_outputs.shape[1]

        # 把解码器隐藏状态复制 src_len 次
        decoder_hidden = decoder_hidden.unsqueeze(1).repeat(1, src_len, 1)

        # 计算能量(energy)
        energy = torch.tanh(
            self.W_enc(encoder_outputs) + self.W_dec(decoder_hidden)
        )

        # 注意力得分
        attention = self.v(energy).squeeze(2)  # (batch, src_len)
        attention_weights = torch.softmax(attention, dim=1)

        # 上下文向量 = 加权求和
        context = torch.bmm(
            attention_weights.unsqueeze(1),  # (batch, 1, src_len)
            encoder_outputs                  # (batch, src_len, hidden*2)
        ).squeeze(1)  # (batch, hidden*2)

        return context, attention_weights

class LuongAttention(nn.Module):
    """Luong (Multiplicative) Attention"""

    def __init__(self, hidden_dim, method='dot'):
        super().__init__()
        self.method = method

        if method == 'general':
            self.W = nn.Linear(hidden_dim, hidden_dim)
        elif method == 'concat':
            self.W = nn.Linear(hidden_dim * 2, hidden_dim)
            self.v = nn.Linear(hidden_dim, 1)

    def forward(self, decoder_hidden, encoder_outputs):
        """
        decoder_hidden: (batch, hidden)
        encoder_outputs: (batch, src_len, hidden)
        """
        if self.method == 'dot':
            # 内积
            score = torch.bmm(
                encoder_outputs,
                decoder_hidden.unsqueeze(2)
            ).squeeze(2)

        elif self.method == 'general':
            # W * h_encoder
            energy = self.W(encoder_outputs)
            score = torch.bmm(
                energy,
                decoder_hidden.unsqueeze(2)
            ).squeeze(2)

        elif self.method == 'concat':
            decoder_expanded = decoder_hidden.unsqueeze(1).expand_as(encoder_outputs)
            energy = torch.tanh(self.W(torch.cat([decoder_expanded, encoder_outputs], dim=2)))
            score = self.v(energy).squeeze(2)

        attention_weights = torch.softmax(score, dim=1)
        context = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs).squeeze(1)

        return context, attention_weights

class SelfAttention(nn.Module):
    """Self-Attention 实现"""

    def __init__(self, embed_dim, num_heads=8, dropout=0.1):
        super().__init__()
        assert embed_dim % num_heads == 0

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5

        # Query, Key, Value 投影
        self.W_q = nn.Linear(embed_dim, embed_dim)
        self.W_k = nn.Linear(embed_dim, embed_dim)
        self.W_v = nn.Linear(embed_dim, embed_dim)
        self.W_o = nn.Linear(embed_dim, embed_dim)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        """
        x: (batch, seq_len, embed_dim)
        mask: (batch, seq_len, seq_len) 注意力掩码
        """
        batch_size, seq_len, _ = x.shape

        # 计算 Query, Key, Value
        Q = self.W_q(x)  # (batch, seq, embed)
        K = self.W_k(x)
        V = self.W_v(x)

        # 为多头变换形状
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        # Q, K, V: (batch, num_heads, seq_len, head_dim)

        # Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale
        # scores: (batch, num_heads, seq_len, seq_len)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attention_weights = torch.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # 应用注意力
        attended = torch.matmul(attention_weights, V)
        # attended: (batch, num_heads, seq_len, head_dim)

        # 合并各个头
        attended = attended.transpose(1, 2).contiguous()
        attended = attended.view(batch_size, seq_len, self.embed_dim)

        # 输出投影
        output = self.W_o(attended)

        return output, attention_weights

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    """正弦/余弦 Positional Encoding"""

    def __init__(self, embed_dim, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 计算 PE 矩阵
        pe = torch.zeros(max_len, embed_dim)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, embed_dim, 2).float() * (-math.log(10000.0) / embed_dim)
        )

        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数索引: sin
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数索引: cos

        pe = pe.unsqueeze(0)  # (1, max_len, embed_dim)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """x: (batch, seq_len, embed_dim)"""
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

class LearnablePositionalEncoding(nn.Module):
    """可学习的 Positional Encoding(BERT 风格)"""

    def __init__(self, embed_dim, max_len=512):
        super().__init__()
        self.pe = nn.Embedding(max_len, embed_dim)

    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0)
        return x + self.pe(positions)

class MultiHeadAttention(nn.Module):
    """Multi-Head Attention"""

    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super().__init__()
        assert embed_dim % num_heads == 0

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.W_q = nn.Linear(embed_dim, embed_dim, bias=False)
        self.W_k = nn.Linear(embed_dim, embed_dim, bias=False)
        self.W_v = nn.Linear(embed_dim, embed_dim, bias=False)
        self.W_o = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)

    def split_heads(self, x):
        """(batch, seq, embed) -> (batch, heads, seq, head_dim)"""
        batch, seq, _ = x.shape
        x = x.view(batch, seq, self.num_heads, self.head_dim)
        return x.transpose(1, 2)

    def forward(self, query, key, value, mask=None):
        batch_size = query.shape[0]

        Q = self.split_heads(self.W_q(query))
        K = self.split_heads(self.W_k(key))
        V = self.split_heads(self.W_v(value))

        # Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attn_weights = self.dropout(torch.softmax(scores, dim=-1))

        output = torch.matmul(attn_weights, V)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)

        return self.W_o(output), attn_weights

class FeedForward(nn.Module):
    """Position-wise Feed-Forward Network"""

    def __init__(self, embed_dim, ff_dim, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.GELU(),  # 论文中用的是 ReLU，BERT/GPT 中主要用 GELU
            nn.Dropout(dropout),
            nn.Linear(ff_dim, embed_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)

class EncoderLayer(nn.Module):
    """Transformer Encoder Layer"""

    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
        super().__init__()
        self.self_attention = MultiHeadAttention(embed_dim, num_heads, dropout)
        self.feed_forward = FeedForward(embed_dim, ff_dim, dropout)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, src_mask=None):
        # Self-Attention + Add & Norm(Pre-LN 风格)
        attn_output, _ = self.self_attention(x, x, x, src_mask)
        x = self.norm1(x + self.dropout(attn_output))

        # Feed-Forward + Add & Norm
        ff_output = self.feed_forward(x)
        x = self.norm2(x + ff_output)

        return x

class DecoderLayer(nn.Module):
    """Transformer Decoder Layer"""

    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
        super().__init__()
        self.self_attention = MultiHeadAttention(embed_dim, num_heads, dropout)
        self.cross_attention = MultiHeadAttention(embed_dim, num_heads, dropout)
        self.feed_forward = FeedForward(embed_dim, ff_dim, dropout)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.norm3 = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, memory, src_mask=None, tgt_mask=None):
        # Masked Self-Attention(对自身的注意力，屏蔽未来词元)
        self_attn_out, _ = self.self_attention(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(self_attn_out))

        # Cross-Attention(对编码器输出的注意力)
        cross_attn_out, cross_attn_weights = self.cross_attention(x, memory, memory, src_mask)
        x = self.norm2(x + self.dropout(cross_attn_out))

        # Feed-Forward
        ff_out = self.feed_forward(x)
        x = self.norm3(x + ff_out)

        return x, cross_attn_weights

class Transformer(nn.Module):
    """完整的 Transformer 实现"""

    def __init__(
        self,
        src_vocab_size,
        tgt_vocab_size,
        embed_dim=512,
        num_heads=8,
        num_encoder_layers=6,
        num_decoder_layers=6,
        ff_dim=2048,
        max_len=5000,
        dropout=0.1
    ):
        super().__init__()

        # 嵌入层
        self.src_embedding = nn.Embedding(src_vocab_size, embed_dim)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, embed_dim)
        self.pos_encoding = PositionalEncoding(embed_dim, max_len, dropout)
        self.embed_scale = math.sqrt(embed_dim)

        # 编码器堆栈
        self.encoder_layers = nn.ModuleList([
            EncoderLayer(embed_dim, num_heads, ff_dim, dropout)
            for _ in range(num_encoder_layers)
        ])

        # 解码器堆栈
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(embed_dim, num_heads, ff_dim, dropout)
            for _ in range(num_decoder_layers)
        ])

        # 输出层
        self.output_norm = nn.LayerNorm(embed_dim)
        self.output_projection = nn.Linear(embed_dim, tgt_vocab_size)

        # 权重初始化
        self._init_weights()

    def _init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def make_causal_mask(self, seq_len, device):
        """自回归掩码(屏蔽未来词元)"""
        mask = torch.tril(torch.ones(seq_len, seq_len, device=device)).bool()
        return mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq, seq)

    def make_pad_mask(self, x, pad_idx=0):
        """填充掩码"""
        return (x != pad_idx).unsqueeze(1).unsqueeze(2)  # (batch, 1, 1, seq)

    def encode(self, src, src_mask=None):
        x = self.pos_encoding(self.src_embedding(src) * self.embed_scale)
        for layer in self.encoder_layers:
            x = layer(x, src_mask)
        return x

    def decode(self, tgt, memory, src_mask=None, tgt_mask=None):
        x = self.pos_encoding(self.tgt_embedding(tgt) * self.embed_scale)
        for layer in self.decoder_layers:
            x, _ = layer(x, memory, src_mask, tgt_mask)
        return self.output_norm(x)

    def forward(self, src, tgt, src_pad_idx=0, tgt_pad_idx=0):
        # 生成掩码
        src_mask = self.make_pad_mask(src, src_pad_idx)
        tgt_len = tgt.shape[1]
        tgt_pad_mask = self.make_pad_mask(tgt, tgt_pad_idx)
        tgt_causal_mask = self.make_causal_mask(tgt_len, tgt.device)
        tgt_mask = tgt_pad_mask & tgt_causal_mask

        # 编码器
        memory = self.encode(src, src_mask)

        # 解码器
        output = self.decode(tgt, memory, src_mask, tgt_mask)

        # 输出
        logits = self.output_projection(output)

        return logits

# 实例化模型并测试
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = Transformer(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    embed_dim=512,
    num_heads=8,
    num_encoder_layers=6,
    num_decoder_layers=6,
    ff_dim=2048,
    dropout=0.1
).to(device)

# 计算参数数量
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"总参数数量: {total_params:,}")

# 测试
src = torch.randint(1, 10000, (4, 20)).to(device)
tgt = torch.randint(1, 10000, (4, 18)).to(device)
output = model(src, tgt)
print(f"输出形状: {output.shape}")  # (4, 18, 10000)

from transformers import (
    BertTokenizer,
    BertModel,
    BertForSequenceClassification,
    BertForTokenClassification,
    BertForQuestionAnswering,
    AdamW,
    get_linear_schedule_with_warmup
)
import torch
from torch.utils.data import Dataset, DataLoader

# 加载分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# 基本编码
text = "Natural language processing is transforming AI."
encoding = tokenizer(
    text,
    add_special_tokens=True,    # 添加 [CLS], [SEP]
    max_length=128,
    padding='max_length',
    truncation=True,
    return_tensors='pt'
)

print("输入 ID:", encoding['input_ids'][0][:10])
print("注意力掩码:", encoding['attention_mask'][0][:10])
print("解码:", tokenizer.decode(encoding['input_ids'][0]))

# 确认 WordPiece 分词
complex_words = ["unbelievable", "preprocessing", "transformers", "tokenization"]
for word in complex_words:
    tokens = tokenizer.tokenize(word)
    print(f"{word:20} -> {tokens}")

class SentimentDataset(Dataset):
    """情感分析数据集"""

    def __init__(self, texts, labels, tokenizer, max_length=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(),
            'attention_mask': encoding['attention_mask'].squeeze(),
            'label': torch.tensor(self.labels[idx], dtype=torch.long)
        }

def train_bert_classifier(
    train_texts,
    train_labels,
    val_texts,
    val_labels,
    num_labels=2,
    epochs=3,
    batch_size=16,
    lr=2e-5
):
    """BERT 情感分析微调"""

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

    # 创建数据集
    train_dataset = SentimentDataset(train_texts, train_labels, tokenizer)
    val_dataset = SentimentDataset(val_texts, val_labels, tokenizer)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)

    # 加载模型
    model = BertForSequenceClassification.from_pretrained(
        'bert-base-uncased',
        num_labels=num_labels
    ).to(device)

    # 优化器
    optimizer = AdamW(model.parameters(), lr=lr, weight_decay=0.01)

    # 调度器
    total_steps = len(train_loader) * epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=total_steps // 10,
        num_training_steps=total_steps
    )

    # 训练循环
    best_val_acc = 0

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        correct = 0

        for batch in train_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['label'].to(device)

            optimizer.zero_grad()
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels
            )

            loss = outputs.loss
            logits = outputs.logits

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()

            total_loss += loss.item()
            preds = logits.argmax(dim=1)
            correct += (preds == labels).sum().item()

        train_acc = correct / len(train_dataset)
        avg_loss = total_loss / len(train_loader)

        # 验证
        model.eval()
        val_correct = 0
        with torch.no_grad():
            for batch in val_loader:
                input_ids = batch['input_ids'].to(device)
                attention_mask = batch['attention_mask'].to(device)
                labels = batch['label'].to(device)

                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                preds = outputs.logits.argmax(dim=1)
                val_correct += (preds == labels).sum().item()

        val_acc = val_correct / len(val_dataset)
        print(f"Epoch {epoch+1}: Loss={avg_loss:.4f}, Train Acc={train_acc:.4f}, Val Acc={val_acc:.4f}")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), 'best_bert_classifier.pt')

    return model

# 情感分析示例
texts = [
    "This movie was absolutely fantastic! I loved every moment.",
    "Terrible film. Complete waste of time and money.",
    "The acting was decent but the plot was confusing.",
    "One of the best movies I've seen this year!",
    "I fell asleep halfway through. So boring."
]
labels = [1, 0, 0, 1, 0]  # 1: positive, 0: negative

from transformers import BertForTokenClassification

# NER 标签
ner_labels = ['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']
label2id = {label: i for i, label in enumerate(ner_labels)}
id2label = {i: label for i, label in enumerate(ner_labels)}

model = BertForTokenClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=len(ner_labels),
    id2label=id2label,
    label2id=label2id
)

def predict_ner(text, model, tokenizer):
    """NER 预测"""
    model.eval()

    encoding = tokenizer(
        text.split(),
        is_split_into_words=True,
        return_offsets_mapping=True,
        padding=True,
        truncation=True,
        return_tensors='pt'
    )

    with torch.no_grad():
        outputs = model(
            input_ids=encoding['input_ids'],
            attention_mask=encoding['attention_mask']
        )

    predictions = outputs.logits.argmax(dim=2).squeeze().tolist()
    word_ids = encoding.word_ids()

    # 处理子词词元
    result = []
    prev_word_id = None
    for pred, word_id in zip(predictions, word_ids):
        if word_id is None or word_id == prev_word_id:
            continue
        word = text.split()[word_id]
        label = id2label[pred]
        result.append((word, label))
        prev_word_id = word_id

    return result

# 使用示例
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
text = "Barack Obama was born in Hawaii and served as the 44th President of the United States."
# 打印 NER 结果(使用训练好的模型时)
# print(predict_ner(text, model, tokenizer))

from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

# 加载 GPT-2
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')
model.eval()

tokenizer.pad_token = tokenizer.eos_token

def generate_text(
    prompt,
    model,
    tokenizer,
    max_length=200,
    temperature=0.9,
    top_p=0.95,
    top_k=50,
    num_return_sequences=1,
    do_sample=True
):
    """GPT-2 文本生成"""

    inputs = tokenizer(prompt, return_tensors='pt')
    input_ids = inputs['input_ids']

    with torch.no_grad():
        output = model.generate(
            input_ids,
            max_length=max_length,
            temperature=temperature,      # 越高输出越多样
            top_p=top_p,                  # nucleus sampling
            top_k=top_k,                  # top-k sampling
            num_return_sequences=num_return_sequences,
            do_sample=do_sample,
            pad_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.2        # 防止重复
        )

    generated_texts = []
    for out in output:
        text = tokenizer.decode(out, skip_special_tokens=True)
        generated_texts.append(text)

    return generated_texts

# 文本生成示例
prompt = "Artificial intelligence is transforming the world by"
generated = generate_text(prompt, model, tokenizer, max_length=150, temperature=0.8)

for i, text in enumerate(generated):
    print(f"\n生成的文本 {i+1}:")
    print(text)
    print("-" * 50)

# Greedy 与 Sampling 比较
def compare_generation_strategies(prompt, model, tokenizer):
    """比较多种生成策略"""

    input_ids = tokenizer.encode(prompt, return_tensors='pt')

    strategies = {
        "Greedy Search": dict(do_sample=False),
        "Beam Search": dict(do_sample=False, num_beams=5, early_stopping=True),
        "Temperature Sampling": dict(do_sample=True, temperature=0.7),
        "Top-k Sampling": dict(do_sample=True, top_k=50),
        "Top-p (Nucleus) Sampling": dict(do_sample=True, top_p=0.92),
    }

    print(f"提示词: {prompt}\n")

    with torch.no_grad():
        for name, params in strategies.items():
            output = model.generate(
                input_ids,
                max_new_tokens=50,
                pad_token_id=tokenizer.eos_token_id,
                **params
            )
            text = tokenizer.decode(output[0], skip_special_tokens=True)
            print(f"[{name}]")
            print(text[len(prompt):].strip())
            print()

compare_generation_strategies(
    "The future of artificial intelligence",
    model,
    tokenizer
)

def zero_shot_classification(text, categories, model_name="gpt2"):
    """Zero-shot 分类"""

    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model = GPT2LMHeadModel.from_pretrained(model_name)
    model.eval()

    # 计算每个类别的对数概率
    scores = {}

    for category in categories:
        prompt = f"Text: {text}\nCategory: {category}"
        inputs = tokenizer(prompt, return_tensors='pt')

        with torch.no_grad():
            outputs = model(**inputs, labels=inputs['input_ids'])
            # 更低的 perplexity = 更有可能的类别
            scores[category] = -outputs.loss.item()

    best_category = max(scores, key=scores.get)
    return best_category, scores

# Few-shot 学习示例
few_shot_prompt = """
Classify the sentiment of the following texts:

Text: "This movie was amazing!"
Sentiment: Positive

Text: "I hated every minute of it."
Sentiment: Negative

Text: "The film was okay, nothing special."
Sentiment: Neutral

Text: "Absolutely brilliant performance!"
Sentiment:"""

print("Few-shot 提示词:")
print(few_shot_prompt)

# Chain of Thought 提示
cot_prompt = """
Q: A train travels 120 miles in 2 hours. How long will it take to travel 300 miles?

A: Let me think step by step.
1. First, find the speed: 120 miles / 2 hours = 60 miles per hour
2. Then, calculate the time for 300 miles: 300 miles / 60 mph = 5 hours
Therefore, it will take 5 hours.

Q: If a store sells 3 apples for $2.40, how much does 7 apples cost?

A: Let me think step by step.
"""

from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def setup_lora_model(model_name="gpt2", r=8, lora_alpha=32, lora_dropout=0.1):
    """LoRA 模型设置"""

    # 加载基础模型
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token

    # LoRA 配置
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=r,                          # LoRA 秩
        lora_alpha=lora_alpha,        # 缩放参数
        target_modules=["c_attn"],    # 要应用的层
        lora_dropout=lora_dropout,
        bias="none"
    )

    # 创建 LoRA 模型
    model = get_peft_model(model, lora_config)

    # 打印可训练参数比例
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"可训练参数: {trainable:,} ({100 * trainable / total:.2f}%)")

    return model, tokenizer

# LoRA 训练示例
class InstructionDataset(torch.utils.data.Dataset):
    """指令格式数据集"""

    def __init__(self, data, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length

        # 转换为指令格式
        self.texts = []
        for item in data:
            prompt = f"### Instruction:\n{item['instruction']}\n\n### Response:\n{item['response']}"
            self.texts.append(prompt)

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        input_ids = encoding['input_ids'].squeeze()
        return {
            'input_ids': input_ids,
            'attention_mask': encoding['attention_mask'].squeeze(),
            'labels': input_ids.clone()  # 语言建模: 输入 = 标签
        }

from transformers import AutoTokenizer, AutoModel
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

class SimpleRAG:
    """简单的 RAG 系统实现"""

    def __init__(self, embedding_model="sentence-transformers/all-MiniLM-L6-v2"):
        self.tokenizer = AutoTokenizer.from_pretrained(embedding_model)
        self.model = AutoModel.from_pretrained(embedding_model)
        self.knowledge_base = []
        self.embeddings = []

    def mean_pooling(self, model_output, attention_mask):
        """用均值池化计算句子嵌入"""
        token_embeddings = model_output[0]
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(
            token_embeddings.size()
        ).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / \
               torch.clamp(input_mask_expanded.sum(1), min=1e-9)

    def encode(self, texts):
        """把文本转换为嵌入向量"""
        encoded = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors='pt'
        )

        with torch.no_grad():
            model_output = self.model(**encoded)

        embeddings = self.mean_pooling(model_output, encoded['attention_mask'])
        # L2 归一化
        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
        return embeddings.numpy()

    def add_documents(self, documents):
        """向知识库添加文档"""
        self.knowledge_base.extend(documents)
        new_embeddings = self.encode(documents)
        if len(self.embeddings) == 0:
            self.embeddings = new_embeddings
        else:
            self.embeddings = np.vstack([self.embeddings, new_embeddings])
        print(f"已添加 {len(documents)} 篇文档。共 {len(self.knowledge_base)} 篇。")

    def retrieve(self, query, top_k=3):
        """检索与查询最相似的文档"""
        query_embedding = self.encode([query])
        similarities = cosine_similarity(query_embedding, self.embeddings)[0]
        top_indices = np.argsort(similarities)[::-1][:top_k]

        results = []
        for idx in top_indices:
            results.append({
                'document': self.knowledge_base[idx],
                'similarity': similarities[idx]
            })

        return results

    def answer(self, query, generator_model, generator_tokenizer, top_k=3):
        """RAG: 检索 + 生成"""
        # 检索相关文档
        relevant_docs = self.retrieve(query, top_k=top_k)

        # 构建上下文
        context = "\n".join([doc['document'] for doc in relevant_docs])

        # 构建提示词
        prompt = f"""Based on the following context, answer the question.

Context:
{context}

Question: {query}

Answer:"""

        # 生成文本
        inputs = generator_tokenizer(prompt, return_tensors='pt', max_length=512, truncation=True)
        with torch.no_grad():
            outputs = generator_model.generate(
                inputs['input_ids'],
                max_new_tokens=150,
                temperature=0.7,
                do_sample=True,
                pad_token_id=generator_tokenizer.eos_token_id
            )

        answer = generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
        answer = answer[len(prompt):].strip()

        return answer, relevant_docs

# RAG 系统使用示例
rag = SimpleRAG()

documents = [
    "BERT (Bidirectional Encoder Representations from Transformers) was published by Google in 2018.",
    "GPT-3 has 175 billion parameters and was developed by OpenAI in 2020.",
    "The Transformer architecture was introduced in the paper 'Attention is All You Need' in 2017.",
    "RLHF (Reinforcement Learning from Human Feedback) is used to align language models with human values.",
    "LoRA allows efficient fine-tuning by adding low-rank matrices to pre-trained model weights.",
    "RAG combines information retrieval with text generation for more accurate responses.",
]

rag.add_documents(documents)

query = "What is BERT and when was it created?"
results = rag.retrieve(query, top_k=2)

print(f"\n查询: {query}")
print("\n检索到的文档:")
for r in results:
    print(f"  [{r['similarity']:.3f}] {r['document']}")