min over theta: E over T ~ p(T) [L_T(f_theta)]

import torch
import torch.nn as nn
import numpy as np
from typing import List, Tuple, Dict


def create_episode(
    dataset,
    n_way: int,
    k_shot: int,
    n_query: int,
    classes: List[int] = None
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    生成 N-way K-shot 情节（episode）
    Returns: (support_x, support_y, query_x, query_y)
    """
    # 选择类别
    all_classes = list(set(dataset.targets.tolist()))
    if classes is None:
        selected_classes = np.random.choice(
            all_classes, n_way, replace=False
        )
    else:
        selected_classes = classes

    support_x, support_y = [], []
    query_x, query_y = [], []

    for new_label, cls in enumerate(selected_classes):
        # 该类别的所有索引
        cls_indices = (dataset.targets == cls).nonzero(as_tuple=True)[0]
        chosen = np.random.choice(
            len(cls_indices),
            k_shot + n_query,
            replace=False
        )

        for i, idx in enumerate(cls_indices[chosen]):
            x, _ = dataset[idx.item()]
            if i < k_shot:
                support_x.append(x)
                support_y.append(new_label)
            else:
                query_x.append(x)
                query_y.append(new_label)

    support_x = torch.stack(support_x)
    support_y = torch.tensor(support_y)
    query_x = torch.stack(query_x)
    query_y = torch.tensor(query_y)

    return support_x, support_y, query_x, query_y

y_hat = sum over i: a(x_hat, x_i) * y_i

class MatchingNetworks(nn.Module):
    """Matching Networks 实现"""

    def __init__(self, encoder: nn.Module, use_fce: bool = False):
        """
        encoder: 特征提取器
        use_fce: 是否使用 Full Context Embedding
        """
        super().__init__()
        self.encoder = encoder
        self.use_fce = use_fce

    def cosine_similarity(
        self,
        query: torch.Tensor,
        support: torch.Tensor
    ) -> torch.Tensor:
        """
        计算余弦相似度
        query: (n_query, embed_dim)
        support: (n_support, embed_dim)
        Returns: (n_query, n_support)
        """
        query_norm = nn.functional.normalize(query, dim=-1)
        support_norm = nn.functional.normalize(support, dim=-1)
        return torch.mm(query_norm, support_norm.t())

    def forward(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        query_x: torch.Tensor,
        n_way: int
    ) -> torch.Tensor:
        """
        support_x: (n_way * k_shot, C, H, W)
        support_y: (n_way * k_shot,)
        query_x: (n_query, C, H, W)
        """
        n_support = support_x.size(0)
        n_query = query_x.size(0)

        # 编码
        support_emb = self.encoder(support_x)   # (n_support, D)
        query_emb = self.encoder(query_x)        # (n_query, D)

        # 计算相似度
        similarities = self.cosine_similarity(
            query_emb, support_emb
        )  # (n_query, n_support)

        # Softmax 注意力
        attention = nn.functional.softmax(similarities, dim=-1)

        # 转换为 one-hot 标签
        support_labels_one_hot = nn.functional.one_hot(
            support_y, n_way
        ).float()  # (n_support, n_way)

        # 以注意力加权和进行预测
        # (n_query, n_support) x (n_support, n_way) = (n_query, n_way)
        logits = torch.mm(attention, support_labels_one_hot)

        return logits  # 以 log 概率形式返回

p_c = (1/|S_c|) sum over (x_i, y_i) in S_c: f_phi(x_i)

p(y=c | x) = softmax(-d(f_phi(x), p_c))

class ConvEncoder(nn.Module):
    """少样本学习用的 4 层 CNN 编码器"""

    def __init__(
        self,
        in_channels: int = 1,
        hidden_dim: int = 64,
        out_dim: int = 64
    ):
        super().__init__()

        def conv_block(in_ch, out_ch):
            return nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=1),
                nn.BatchNorm2d(out_ch),
                nn.ReLU(),
                nn.MaxPool2d(2)
            )

        self.net = nn.Sequential(
            conv_block(in_channels, hidden_dim),
            conv_block(hidden_dim, hidden_dim),
            conv_block(hidden_dim, hidden_dim),
            conv_block(hidden_dim, out_dim),
            nn.Flatten()
        )

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


class PrototypicalNetworks(nn.Module):
    """Prototypical Networks 完整实现"""

    def __init__(self, encoder: nn.Module):
        super().__init__()
        self.encoder = encoder

    def compute_prototypes(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        n_way: int
    ) -> torch.Tensor:
        """
        计算各类别的原型（嵌入均值）
        support_x: (n_way * k_shot, C, H, W)
        support_y: (n_way * k_shot,)
        Returns: (n_way, embed_dim)
        """
        support_emb = self.encoder(support_x)  # (n_support, D)
        prototypes = []

        for cls in range(n_way):
            mask = (support_y == cls)
            cls_embeddings = support_emb[mask]
            prototype = cls_embeddings.mean(dim=0)
            prototypes.append(prototype)

        return torch.stack(prototypes)  # (n_way, D)

    def euclidean_dist(
        self,
        x: torch.Tensor,
        y: torch.Tensor
    ) -> torch.Tensor:
        """
        计算欧几里得距离
        x: (n, D)
        y: (m, D)
        Returns: (n, m)
        """
        # ||x - y||^2 = ||x||^2 + ||y||^2 - 2*x·y
        n = x.size(0)
        m = y.size(0)
        x_sq = (x ** 2).sum(dim=1, keepdim=True).expand(n, m)
        y_sq = (y ** 2).sum(dim=1, keepdim=True).expand(m, n).t()
        xy = torch.mm(x, y.t())
        dist = x_sq + y_sq - 2 * xy
        return dist.clamp(min=0).sqrt()

    def forward(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        query_x: torch.Tensor,
        n_way: int
    ) -> torch.Tensor:
        """
        Prototypical Networks 前向传播
        Returns: 查询样本的 log 概率 (n_query, n_way)
        """
        # 计算原型
        prototypes = self.compute_prototypes(
            support_x, support_y, n_way
        )  # (n_way, D)

        # 查询嵌入
        query_emb = self.encoder(query_x)  # (n_query, D)

        # 计算距离
        dists = self.euclidean_dist(
            query_emb, prototypes
        )  # (n_query, n_way)

        # 用负距离作为 logit（距离越近，概率越高）
        log_probs = nn.functional.log_softmax(-dists, dim=-1)

        return log_probs


# ========== 训练循环 ==========

def train_prototypical(
    model: PrototypicalNetworks,
    train_dataset,
    n_way: int = 5,
    k_shot: int = 5,
    n_query: int = 15,
    n_episodes: int = 100,
    lr: float = 1e-3,
    device: str = 'cpu'
):
    """训练 Prototypical Networks"""
    model = model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.NLLLoss()

    model.train()
    episode_losses = []
    episode_accs = []

    for episode in range(n_episodes):
        # 生成 episode
        support_x, support_y, query_x, query_y = create_episode(
            train_dataset, n_way, k_shot, n_query
        )
        support_x = support_x.to(device)
        support_y = support_y.to(device)
        query_x = query_x.to(device)
        query_y = query_y.to(device)

        optimizer.zero_grad()

        log_probs = model(support_x, support_y, query_x, n_way)
        loss = criterion(log_probs, query_y)
        loss.backward()
        optimizer.step()

        # 计算准确率
        preds = log_probs.argmax(dim=-1)
        acc = (preds == query_y).float().mean().item()

        episode_losses.append(loss.item())
        episode_accs.append(acc)

        if (episode + 1) % 100 == 0:
            print(
                f"Episode {episode+1}/{n_episodes} | "
                f"Loss: {np.mean(episode_losses[-100:]):.4f} | "
                f"Acc: {np.mean(episode_accs[-100:]):.4f}"
            )

    return episode_losses, episode_accs

class RelationNetwork(nn.Module):
    """Relation Networks：可学习的距离函数"""

    def __init__(self, encoder: nn.Module, embed_dim: int = 64):
        super().__init__()
        self.encoder = encoder

        # 关系模块：连接两个嵌入后输出关系得分
        self.relation_module = nn.Sequential(
            nn.Linear(embed_dim * 2, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

    def forward(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        query_x: torch.Tensor,
        n_way: int
    ) -> torch.Tensor:
        """
        Returns: 关系得分 (n_query, n_way)
        """
        support_emb = self.encoder(support_x)
        query_emb = self.encoder(query_x)

        # 各类别的原型
        prototypes = []
        for cls in range(n_way):
            mask = (support_y == cls)
            proto = support_emb[mask].mean(dim=0)
            prototypes.append(proto)
        prototypes = torch.stack(prototypes)  # (n_way, D)

        n_query = query_emb.size(0)

        # 每个查询与每个原型配对的关系得分
        query_expanded = query_emb.unsqueeze(1).expand(
            n_query, n_way, -1
        )  # (n_query, n_way, D)
        proto_expanded = prototypes.unsqueeze(0).expand(
            n_query, n_way, -1
        )  # (n_query, n_way, D)

        # 连接
        pairs = torch.cat(
            [query_expanded, proto_expanded], dim=-1
        )  # (n_query, n_way, 2D)

        # 计算关系得分
        scores = self.relation_module(
            pairs.view(-1, pairs.size(-1))
        ).view(n_query, n_way)

        return scores

theta_i' = theta - alpha * grad_theta L_{T_i}(f_theta)

theta = theta - beta * grad_theta sum_i L_{T_i}(f_{theta_i'})

grad_theta L(f_{theta_i'}) = grad_theta L(f_{theta - alpha * grad L(theta)})

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np
from copy import deepcopy
from typing import List, Tuple


class MAML:
    """
    MAML: Model-Agnostic Meta-Learning
    Finn et al., 2017 (arXiv:1703.03400)
    """

    def __init__(
        self,
        model: nn.Module,
        inner_lr: float = 0.01,       # alpha：内循环学习率
        outer_lr: float = 0.001,      # beta：外循环学习率
        n_inner_steps: int = 5,       # 内循环更新次数
        first_order: bool = False,    # 是否使用 FOMAML
        device: str = 'cpu'
    ):
        self.model = model.to(device)
        self.inner_lr = inner_lr
        self.outer_lr = outer_lr
        self.n_inner_steps = n_inner_steps
        self.first_order = first_order
        self.device = device

        self.meta_optimizer = torch.optim.Adam(
            self.model.parameters(), lr=outer_lr
        )

    def inner_loop(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        model_params=None
    ) -> dict:
        """
        内循环：在 support 集上做任务特定适应
        Returns: 适应后的参数字典
        """
        # 复制当前参数
        if model_params is None:
            params = {
                name: param.clone()
                for name, param in self.model.named_parameters()
            }
        else:
            params = {k: v.clone() for k, v in model_params.items()}

        for step in range(self.n_inner_steps):
            # 用当前参数前向传播（使用 functional API）
            logits = self._forward_with_params(support_x, params)
            loss = F.cross_entropy(logits, support_y)

            # 对参数求梯度
            grads = torch.autograd.grad(
                loss,
                params.values(),
                create_graph=not self.first_order  # 为二阶导数保留计算图
            )

            # 更新参数（SGD）
            params = {
                name: param - self.inner_lr * grad
                for (name, param), grad in zip(params.items(), grads)
            }

        return params

    def _forward_with_params(
        self,
        x: torch.Tensor,
        params: dict
    ) -> torch.Tensor:
        """
        用指定参数进行前向传播
        使用 torch.nn.utils.stateless 或 functorch
        """
        # 临时替换模型的当前参数
        original_params = {}
        for name, param in self.model.named_parameters():
            original_params[name] = param.data
            param.data = params[name].data if name in params else param.data

        output = self.model(x)

        # 恢复原始参数（create_graph=True 时不需要）
        for name, param in self.model.named_parameters():
            if name in original_params:
                param.data = original_params[name]

        return output

    def meta_train_step(
        self,
        tasks: List[Tuple]
    ) -> float:
        """
        一次元训练步骤（对多个任务）
        tasks: [(support_x, support_y, query_x, query_y), ...]
        """
        self.meta_optimizer.zero_grad()

        meta_loss = 0.0

        for support_x, support_y, query_x, query_y in tasks:
            support_x = support_x.to(self.device)
            support_y = support_y.to(self.device)
            query_x = query_x.to(self.device)
            query_y = query_y.to(self.device)

            # 内循环：任务特定适应
            adapted_params = self.inner_loop(support_x, support_y)

            # 外循环：用适应后的参数在 query 集上评估
            query_logits = self._forward_with_params(
                query_x, adapted_params
            )
            query_loss = F.cross_entropy(query_logits, query_y)
            meta_loss += query_loss

        # 按任务数取平均
        meta_loss /= len(tasks)

        # 更新元参数
        meta_loss.backward()
        self.meta_optimizer.step()

        return meta_loss.item()

    def fine_tune(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        n_steps: int = None
    ) -> nn.Module:
        """
        对新任务进行微调（在推理时使用）
        """
        n_steps = n_steps or self.n_inner_steps
        model_copy = deepcopy(self.model)
        optimizer = torch.optim.SGD(
            model_copy.parameters(), lr=self.inner_lr
        )

        model_copy.train()
        for step in range(n_steps):
            optimizer.zero_grad()
            logits = model_copy(support_x.to(self.device))
            loss = F.cross_entropy(logits, support_y.to(self.device))
            loss.backward()
            optimizer.step()

        return model_copy

    def evaluate(
        self,
        tasks: List[Tuple],
        n_fine_tune_steps: int = 5
    ) -> Tuple[float, float]:
        """
        元测试：对新任务适应后进行评估
        """
        total_loss = 0.0
        total_acc = 0.0

        for support_x, support_y, query_x, query_y in tasks:
            # 微调
            adapted_model = self.fine_tune(support_x, support_y, n_fine_tune_steps)
            adapted_model.eval()

            with torch.no_grad():
                query_logits = adapted_model(query_x.to(self.device))
                loss = F.cross_entropy(query_logits, query_y.to(self.device))
                preds = query_logits.argmax(dim=-1)
                acc = (preds == query_y.to(self.device)).float().mean()

            total_loss += loss.item()
            total_acc += acc.item()

        return total_loss / len(tasks), total_acc / len(tasks)


# ========== MAML 训练循环 ==========

def train_maml(
    maml: MAML,
    dataset,
    n_way: int = 5,
    k_shot: int = 1,
    n_query: int = 15,
    meta_batch_size: int = 32,
    n_iterations: int = 60000,
    device: str = 'cpu'
):
    """MAML 训练"""
    print(f"MAML 训练开始：{n_way}-way {k_shot}-shot")
    print(f"元批次大小：{meta_batch_size}")
    print(f"总迭代次数：{n_iterations}")

    losses = []

    for iteration in range(n_iterations):
        # 生成元批次
        tasks = []
        for _ in range(meta_batch_size):
            task = create_episode(dataset, n_way, k_shot, n_query)
            tasks.append(task)

        # 元训练步骤
        meta_loss = maml.meta_train_step(tasks)
        losses.append(meta_loss)

        if (iteration + 1) % 1000 == 0:
            avg_loss = np.mean(losses[-1000:])
            print(f"迭代 {iteration+1}/{n_iterations} | 元损失: {avg_loss:.4f}")

    return losses

theta = theta + epsilon * (W_k - theta)

class Reptile:
    """
    Reptile: A Scalable Meta-learning Algorithm
    Nichol et al., 2018 (arXiv:1803.02999)
    """

    def __init__(
        self,
        model: nn.Module,
        inner_lr: float = 0.02,
        outer_lr: float = 0.001,
        n_inner_steps: int = 5,
        device: str = 'cpu'
    ):
        self.model = model.to(device)
        self.inner_lr = inner_lr
        self.outer_lr = outer_lr  # epsilon
        self.n_inner_steps = n_inner_steps
        self.device = device

    def inner_train(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor
    ) -> dict:
        """任务内循环训练（SGD k 次）"""
        model_copy = deepcopy(self.model)
        optimizer = torch.optim.SGD(
            model_copy.parameters(), lr=self.inner_lr
        )

        model_copy.train()
        for step in range(self.n_inner_steps):
            optimizer.zero_grad()
            logits = model_copy(support_x)
            loss = F.cross_entropy(logits, support_y)
            loss.backward()
            optimizer.step()

        return dict(model_copy.named_parameters())

    def meta_update(self, task_params_list: List[dict]):
        """
        Reptile 元更新：
        theta += epsilon * (mean(W_k) - theta)
        """
        with torch.no_grad():
            for name, param in self.model.named_parameters():
                # 各任务参数的平均值
                task_mean = torch.stack([
                    task_params[name].data
                    for task_params in task_params_list
                ]).mean(dim=0)

                # Reptile 更新
                param.data += self.outer_lr * (task_mean - param.data)

    def train(
        self,
        dataset,
        n_way: int = 5,
        k_shot: int = 5,
        meta_batch_size: int = 5,
        n_iterations: int = 100000
    ):
        """Reptile 完整训练循环"""
        print(f"Reptile 训练开始：{n_way}-way {k_shot}-shot")

        for iteration in range(n_iterations):
            task_params_list = []

            for _ in range(meta_batch_size):
                support_x, support_y, _, _ = create_episode(
                    dataset, n_way, k_shot, n_query=0
                )
                support_x = support_x.to(self.device)
                support_y = support_y.to(self.device)

                task_params = self.inner_train(support_x, support_y)
                task_params_list.append(task_params)

            self.meta_update(task_params_list)

            if (iteration + 1) % 10000 == 0:
                print(f"迭代 {iteration+1}/{n_iterations} 完成")

英语→法语翻译：
sea otter => loutre de mer
peppermint => menthe poivrée
plush giraffe => girafe en peluche
cheese => ?

from typing import List, Dict, Any
import numpy as np


class FewShotPromptBuilder:
    """少样本提示词构建器"""

    def __init__(self):
        self.examples = []
        self.instruction = ""
        self.template = "{input} => {output}"

    def set_instruction(self, instruction: str):
        """设置任务指令"""
        self.instruction = instruction
        return self

    def add_example(self, input_text: str, output_text: str):
        """添加示例"""
        self.examples.append({
            'input': input_text,
            'output': output_text
        })
        return self

    def build_prompt(self, query: str) -> str:
        """构建完整的少样本提示词"""
        parts = []

        if self.instruction:
            parts.append(self.instruction)
            parts.append("")

        for ex in self.examples:
            parts.append(self.template.format(
                input=ex['input'],
                output=ex['output']
            ))

        # 查询（输出留空）
        parts.append(f"{query} =>")

        return "\n".join(parts)

    def build_chat_messages(
        self,
        query: str,
        system_prompt: str = None
    ) -> List[Dict]:
        """构建 Chat 格式消息（用于 GPT-4、Claude 等）"""
        messages = []

        if system_prompt:
            messages.append({
                'role': 'system',
                'content': system_prompt
            })

        # 将示例转换为对话形式
        for ex in self.examples:
            messages.append({
                'role': 'user',
                'content': ex['input']
            })
            messages.append({
                'role': 'assistant',
                'content': ex['output']
            })

        # 实际查询
        messages.append({
            'role': 'user',
            'content': query
        })

        return messages


class DynamicExampleSelector:
    """
    动态示例选择器
    选择与查询最相似的示例，实现更有效的少样本学习
    """

    def __init__(self, examples: List[Dict], encoder=None):
        self.examples = examples
        self.encoder = encoder  # sentence-transformers 等

    def select_similar(
        self,
        query: str,
        n_examples: int = 3
    ) -> List[Dict]:
        """
        选择与查询最相似的 n 个示例
        """
        if self.encoder is None:
            # 没有编码器时随机选择
            return np.random.choice(
                self.examples, n_examples, replace=False
            ).tolist()

        # 基于语义相似度选择
        query_emb = self.encoder.encode(query)
        example_embs = self.encoder.encode(
            [ex['input'] for ex in self.examples]
        )

        # 余弦相似度
        similarities = np.dot(example_embs, query_emb) / (
            np.linalg.norm(example_embs, axis=1)
            * np.linalg.norm(query_emb)
        )

        top_indices = np.argsort(similarities)[-n_examples:][::-1]
        return [self.examples[i] for i in top_indices]

    def select_diverse(
        self,
        n_examples: int = 3
    ) -> List[Dict]:
        """
        选择使多样性最大化的示例（MMR 算法）
        """
        if len(self.examples) <= n_examples:
            return self.examples

        if self.encoder is None:
            return np.random.choice(
                self.examples, n_examples, replace=False
            ).tolist()

        embeddings = self.encoder.encode(
            [ex['input'] for ex in self.examples]
        )

        selected = [0]  # 选择第一个示例
        remaining = list(range(1, len(self.examples)))

        while len(selected) < n_examples:
            # 与已选示例的最大相似度
            selected_embs = embeddings[selected]
            best_idx = None
            best_score = float('-inf')

            for idx in remaining:
                # MMR：与查询的相似度 - 与已选示例的相似度
                sim_to_selected = np.max(
                    np.dot(selected_embs, embeddings[idx]) / (
                        np.linalg.norm(selected_embs, axis=1)
                        * np.linalg.norm(embeddings[idx])
                    )
                )
                score = -sim_to_selected  # 最大化多样性

                if score > best_score:
                    best_score = score
                    best_idx = idx

            selected.append(best_idx)
            remaining.remove(best_idx)

        return [self.examples[i] for i in selected]


# ========== 实战示例 ==========

def sentiment_analysis_few_shot():
    """情感分析少样本示例"""
    builder = FewShotPromptBuilder()

    builder.set_instruction("请分析以下电影评论的情感倾向，回答 Positive（正面）或 Negative（负面）。")

    builder.add_example(
        "这部电影真的很感人，演员的演技也非常出色。",
        "Positive"
    )
    builder.add_example(
        "剧情太无聊了，结局也令人失望。",
        "Negative"
    )
    builder.add_example(
        "特效很棒，但剧本太单薄了。",
        "Negative"
    )
    builder.add_example(
        "久违地和家人一起看了一部温馨的电影。",
        "Positive"
    )

    query = "导演风格独特，音乐与画面完美融合。"
    prompt = builder.build_prompt(query)

    print("=== 少样本提示词 ===")
    print(prompt)
    return prompt


def code_generation_few_shot():
    """代码生成少样本示例"""
    builder = FewShotPromptBuilder()

    builder.set_instruction(
        "请将自然语言描述转换为 Python 代码。"
    )

    builder.add_example(
        "找出列表中的最大元素",
        "def find_max(lst):\n    return max(lst)"
    )
    builder.add_example(
        "判断字符串是否为回文",
        "def is_palindrome(s):\n    return s == s[::-1]"
    )
    builder.add_example(
        "展平嵌套列表",
        "def flatten(lst):\n    return [x for sublist in lst for x in sublist]"
    )

    query = "统计列表中每个元素出现的频率"
    prompt = builder.build_prompt(query)

    print("=== 代码生成少样本提示词 ===")
    print(prompt)
    return prompt

class CrossLingualFewShot:
    """
    跨语言少样本学习
    英语示例 + 目标语言查询
    """

    def __init__(self, model_name: str = "xlm-roberta-large"):
        from transformers import AutoTokenizer, AutoModelForSequenceClassification
        import torch

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name, num_labels=2
        )

    def encode_text(self, text: str) -> torch.Tensor:
        """将文本编码为嵌入向量"""
        inputs = self.tokenizer(
            text,
            return_tensors='pt',
            padding=True,
            truncation=True,
            max_length=128
        )
        with torch.no_grad():
            outputs = self.model(**inputs, output_hidden_states=True)
            # [CLS] token 的最后一层隐藏状态
            embedding = outputs.hidden_states[-1][:, 0, :]
        return embedding

    def classify_zero_shot(
        self,
        query_ko: str,
        class_descriptions_en: List[str]
    ) -> int:
        """
        用英语类别描述对韩语查询进行分类
        """
        query_emb = self.encode_text(query_ko)
        class_embs = torch.cat([
            self.encode_text(desc) for desc in class_descriptions_en
        ])

        # 余弦相似度
        similarities = F.cosine_similarity(
            query_emb.expand(len(class_descriptions_en), -1),
            class_embs
        )

        return similarities.argmax().item()

    def few_shot_classify(
        self,
        support_texts: List[str],
        support_labels: List[int],
        query_text: str,
        n_classes: int
    ) -> int:
        """
        少样本分类（原型方法）
        support 与 query 可以使用不同语言
        """
        support_embs = torch.cat([
            self.encode_text(t) for t in support_texts
        ])
        query_emb = self.encode_text(query_text)

        # 各类别的原型
        prototypes = []
        for cls in range(n_classes):
            mask = torch.tensor([l == cls for l in support_labels])
            proto = support_embs[mask].mean(dim=0, keepdim=True)
            prototypes.append(proto)
        prototypes = torch.cat(prototypes)

        # 用欧几里得距离分类
        dists = torch.cdist(query_emb, prototypes)
        return dists.argmin().item()

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np


class MedicalImageEncoder(nn.Module):
    """
    医疗影像编码器
    基于 ResNet-18，针对医疗影像特性进行调整
    """

    def __init__(self, embed_dim: int = 512, pretrained: bool = True):
        super().__init__()

        # ResNet-18 主干网络
        backbone = models.resnet18(pretrained=pretrained)

        # 移除 FC 层
        self.backbone = nn.Sequential(*list(backbone.children())[:-1])

        # 嵌入头
        self.embed_head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512, embed_dim),
            nn.LayerNorm(embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, embed_dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        features = self.backbone(x)
        return self.embed_head(features)


class MedicalFewShotClassifier:
    """
    医疗影像少样本分类器
    基于 Prototypical Networks
    """

    def __init__(
        self,
        encoder: MedicalImageEncoder,
        device: str = 'cpu'
    ):
        self.encoder = encoder.to(device)
        self.device = device
        self.prototypes = {}
        self.disease_names = {}

    def register_disease(
        self,
        disease_id: int,
        disease_name: str,
        support_images: List,
        transform=None
    ):
        """
        注册新疾病（使用少量示例）
        """
        self.disease_names[disease_id] = disease_name
        self.encoder.eval()

        embeddings = []
        with torch.no_grad():
            for img in support_images:
                if transform:
                    img_tensor = transform(img).unsqueeze(0).to(self.device)
                else:
                    img_tensor = img.unsqueeze(0).to(self.device)
                emb = self.encoder(img_tensor)
                embeddings.append(emb)

        prototype = torch.cat(embeddings).mean(dim=0)
        self.prototypes[disease_id] = prototype

        print(
            f"疾病注册完成：{disease_name} "
            f"（{len(support_images)} 个示例）"
        )

    def diagnose(
        self,
        query_image: torch.Tensor,
        top_k: int = 3
    ) -> List[Dict]:
        """
        对查询影像进行诊断
        Returns: 相似度最高的 top-k 个疾病及其得分
        """
        self.encoder.eval()

        with torch.no_grad():
            query_emb = self.encoder(
                query_image.unsqueeze(0).to(self.device)
            )

        results = []
        for disease_id, prototype in self.prototypes.items():
            # 余弦相似度
            similarity = F.cosine_similarity(
                query_emb, prototype.unsqueeze(0)
            ).item()
            results.append({
                'disease_id': disease_id,
                'disease_name': self.disease_names[disease_id],
                'similarity': similarity
            })

        # 按相似度排序
        results.sort(key=lambda x: x['similarity'], reverse=True)
        return results[:top_k]

    def update_prototype(
        self,
        disease_id: int,
        new_image: torch.Tensor,
        momentum: float = 0.9
    ):
        """
        用新的确诊病例更新原型（在线学习）
        """
        self.encoder.eval()

        with torch.no_grad():
            new_emb = self.encoder(
                new_image.unsqueeze(0).to(self.device)
            ).squeeze(0)

        if disease_id in self.prototypes:
            # 用指数移动平均更新
            self.prototypes[disease_id] = (
                momentum * self.prototypes[disease_id]
                + (1 - momentum) * new_emb
            )
            print(f"原型更新完成：{self.disease_names[disease_id]}")
        else:
            print(f"警告：疾病 {disease_id} 尚未注册。")


def demo_medical_few_shot():
    """医疗少样本分类演示"""
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    encoder = MedicalImageEncoder(embed_dim=512, pretrained=True)
    classifier = MedicalFewShotClassifier(encoder)

    print("=== 医疗影像少样本分类系统 ===")
    print("该系统仅凭少量确诊病例即可识别新的疾病。")
    print()

    # 实际使用示例
    # 1. 为每种疾病注册 5~10 张参考影像
    # 2. 输入新患者的影像，即可诊断出相似的疾病
    # 3. 每当出现新的确诊病例，就更新对应原型

    print("使用方法：")
    print("1. classifier.register_disease(id, name, support_images)")
    print("2. results = classifier.diagnose(patient_image)")
    print("3. classifier.update_prototype(disease_id, new_confirmed_image)")

pip install learn2learn

import learn2learn as l2l
import torch
import torch.nn as nn
from typing import Tuple


def build_l2l_maml(
    model: nn.Module,
    lr: float = 0.01,
    first_order: bool = False
) -> l2l.algorithms.MAML:
    """创建 learn2learn MAML 包装器"""
    return l2l.algorithms.MAML(
        model,
        lr=lr,
        first_order=first_order,
        allow_unused=True
    )


def train_with_l2l(
    maml_model: l2l.algorithms.MAML,
    tasksets,
    n_way: int = 5,
    k_shot: int = 1,
    n_query: int = 15,
    meta_lr: float = 0.003,
    n_iterations: int = 1000,
    adaptation_steps: int = 1,
    device: str = 'cpu'
):
    """
    使用 learn2learn 训练 MAML
    """
    maml_model = maml_model.to(device)
    meta_optimizer = torch.optim.Adam(
        maml_model.parameters(), lr=meta_lr
    )
    criterion = nn.CrossEntropyLoss(reduction='mean')

    for iteration in range(n_iterations):
        meta_optimizer.zero_grad()
        meta_train_loss = 0.0
        meta_train_acc = 0.0

        # 迷你元批次
        for task in range(4):  # 同时处理 4 个任务
            # 任务采样
            X, y = tasksets.train.sample()
            X, y = X.to(device), y.to(device)

            # 生成分类器克隆（用于元学习的副本）
            learner = maml_model.clone()

            # 划分 Support / Query
            support_indices = torch.zeros(X.size(0), dtype=torch.bool)
            for cls in range(n_way):
                cls_idx = (y == cls).nonzero(as_tuple=True)[0]
                support_idx = cls_idx[:k_shot]
                support_indices[support_idx] = True

            query_indices = ~support_indices
            support_x, support_y = X[support_indices], y[support_indices]
            query_x, query_y = X[query_indices], y[query_indices]

            # 内循环：adaptation
            for step in range(adaptation_steps):
                support_logits = learner(support_x)
                support_loss = criterion(support_logits, support_y)
                learner.adapt(support_loss)

            # 外循环：meta-gradient
            query_logits = learner(query_x)
            query_loss = criterion(query_logits, query_y)
            meta_train_loss += query_loss

            # 准确率
            preds = query_logits.argmax(dim=-1)
            acc = (preds == query_y).float().mean()
            meta_train_acc += acc

        meta_train_loss /= 4
        meta_train_acc /= 4

        meta_train_loss.backward()
        meta_optimizer.step()

        if (iteration + 1) % 100 == 0:
            print(
                f"迭代 {iteration+1}/{n_iterations} | "
                f"元损失: {meta_train_loss.item():.4f} | "
                f"元准确率: {meta_train_acc.item():.4f}"
            )


def setup_omniglot_maml():
    """
    用 Omniglot 数据集设置 MAML
    Omniglot：50 种字母体系的文字（共 1623 个类别，每类 20 个样本）
    """
    # 使用 learn2learn 的基准数据集
    tasksets = l2l.vision.benchmarks.get_tasksets(
        'omniglot',
        train_ways=5,
        train_samples=2 * 1 + 2 * 15,  # k_shot + n_query
        test_ways=5,
        test_samples=2 * 1 + 2 * 15,
        root='./data',
        device='cpu'
    )

    # CNN 模型
    model = l2l.vision.models.OmniglotCNN(
        output_size=5,
        hidden_size=64,
        layers=4
    )

    # MAML 包装器
    maml = build_l2l_maml(model, lr=0.4, first_order=False)

    return maml, tasksets


def evaluate_l2l(
    maml_model: l2l.algorithms.MAML,
    tasksets,
    n_way: int = 5,
    k_shot: int = 1,
    n_query: int = 15,
    n_test_tasks: int = 600,
    adaptation_steps: int = 3,
    device: str = 'cpu'
) -> Tuple[float, float]:
    """learn2learn 元评估"""
    criterion = nn.CrossEntropyLoss()
    total_loss = 0.0
    total_acc = 0.0

    maml_model.eval()

    for _ in range(n_test_tasks):
        X, y = tasksets.test.sample()
        X, y = X.to(device), y.to(device)

        learner = maml_model.clone()

        support_indices = torch.zeros(X.size(0), dtype=torch.bool)
        for cls in range(n_way):
            cls_idx = (y == cls).nonzero(as_tuple=True)[0]
            support_idx = cls_idx[:k_shot]
            support_indices[support_idx] = True

        query_indices = ~support_indices
        support_x, support_y = X[support_indices], y[support_indices]
        query_x, query_y = X[query_indices], y[query_indices]

        # 测试时使用更多适应步数
        for step in range(adaptation_steps):
            support_loss = criterion(learner(support_x), support_y)
            learner.adapt(support_loss)

        with torch.no_grad():
            query_logits = learner(query_x)
            loss = criterion(query_logits, query_y)
            acc = (query_logits.argmax(dim=-1) == query_y).float().mean()

        total_loss += loss.item()
        total_acc += acc.item()

    return total_loss / n_test_tasks, total_acc / n_test_tasks

def standard_few_shot_evaluation(
    model,
    test_dataset,
    n_way: int = 5,
    k_shot: int = 1,
    n_query: int = 15,
    n_episodes: int = 600,
    confidence_interval: bool = True
) -> Dict:
    """
    标准少样本评测协议
    600 个 episode 的均值与 95% 置信区间
    """
    accs = []
    model.eval()

    for episode in range(n_episodes):
        support_x, support_y, query_x, query_y = create_episode(
            test_dataset, n_way, k_shot, n_query
        )

        with torch.no_grad():
            log_probs = model(support_x, support_y, query_x, n_way)
            preds = log_probs.argmax(dim=-1)
            acc = (preds == query_y).float().mean().item()

        accs.append(acc)

    mean_acc = np.mean(accs)
    std_acc = np.std(accs)

    if confidence_interval:
        # 95% 置信区间
        ci = 1.96 * std_acc / np.sqrt(n_episodes)
        return {
            'mean_accuracy': mean_acc,
            'std': std_acc,
            'confidence_interval_95': ci,
            'result_string': f"{mean_acc*100:.2f} ± {ci*100:.2f}%"
        }

    return {'mean_accuracy': mean_acc, 'std': std_acc}