import torch import torch.nn as nn from typing import Optional, List, Union, Tuple from transformers import LlamaForCausalLM class KoPA(nn.Module): def __init__( self, model: LlamaForCausalLM ) -> None: super(KoPA, self).__init__() self.llama_model = model self.embeddings = nn.Embedding(100, 3072) # self.embeddings = PrefixKGEmbedding( # num_ent=2034, # num_rel=42, # dim_llm=3072, # num_prefix=1 # ) def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, embedding_ids: torch.LongTensor = None ): if embedding_ids.max() >= self.embeddings.num_embeddings or embedding_ids.min() < 0: print(f"[ERROR] embedding_ids 超出范围！最大值: {embedding_ids.max()}, 最小值: {embedding_ids.min()}") embedding_ids = torch.clamp(embedding_ids, min=0, max=self.embeddings.num_embeddings - 1) kg_embeds = self.embeddings(embedding_ids) batch_size, seq_len, _ = kg_embeds.shape token_embeds = self.llama_model.model.model.embed_tokens(input_ids) input_embeds = torch.cat((kg_embeds, token_embeds), dim=1) prefix_mask = torch.ones((batch_size, seq_len)) prefix_labels = torch.full((batch_size, seq_len), fill_value=-100, dtype=torch.long) new_attention_mask = torch.cat((prefix_mask.cuda(), attention_mask), dim=-1) new_labels = torch.cat((prefix_labels.cuda(), labels), dim=-1) if embedding_ids.max() >= self.embeddings.num_embeddings or embedding_ids.min() < 0: print(f"[ERROR] embedding_ids 超出范围！最大值: {embedding_ids.max()}, 最小值: {embedding_ids.min()}") embedding_ids = torch.clamp(embedding_ids, min=0, max=self.embeddings.num_embeddings - 1) return self.llama_model( input_ids=None, attention_mask=new_attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=input_embeds, labels=new_labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) class KoPAWithAdapter(nn.Module): def __init__(self, model, num_prefix, tokenizer=None): super().__init__() self.model = model self.num_prefix = num_prefix hidden_size = model.config.hidden_size # 使用tokenizer获取vocab_size vocab_size = tokenizer.vocab_size if tokenizer else 32000 self.static_prefix_embedding = nn.Embedding(vocab_size, hidden_size) self.embeddings = self.static_prefix_embedding # 保留这个属性 self.sensor_mlp = nn.Sequential( nn.Linear(3, hidden_size // 2), nn.ReLU(), nn.Dropout(0.1), nn.Linear(hidden_size // 2, hidden_size) ) # 添加LayerNorm self.norm = nn.LayerNorm(hidden_size) print(f"[INFO] 模型初始化: hidden_size={hidden_size}, vocab_size={vocab_size}") def forward(self, input_ids, attention_mask, static_prefix=None, sensor_data=None, labels=None, **kwargs): batch_size, seq_len = input_ids.shape device = input_ids.device # 确保所有组件在同一设备上 self.static_prefix_embedding = self.static_prefix_embedding.to(device) self.sensor_mlp = self.sensor_mlp.to(device) self.norm = self.norm.to(device) # 处理静态前缀 if static_prefix is not None: static_prefix = static_prefix.to(device) static_prefix = self.static_prefix_embedding(static_prefix) else: static_prefix = torch.zeros( (batch_size, self.num_prefix, self.model.config.hidden_size), device=device ) # 处理动态前缀 if sensor_data is not None: sensor_data = sensor_data.to(device) if sensor_data.dim() == 1: sensor_data = sensor_data.unsqueeze(0) try: dynamic_prefix = self.sensor_mlp(sensor_data) dynamic_prefix = dynamic_prefix.unsqueeze(1).expand(-1, self.num_prefix, -1) except Exception as e: print(f"[ERROR] sensor_mlp处理失败: {e}") dynamic_prefix = torch.zeros_like(static_prefix) else: dynamic_prefix = torch.zeros_like(static_prefix) # 混合前缀 alpha = 0.6 final_prefix = alpha * static_prefix + (1 - alpha) * dynamic_prefix final_prefix = self.norm(final_prefix) # 处理token嵌入 token_embeds = self.model.model.embed_tokens(input_ids) input_embeds = torch.cat((final_prefix, token_embeds), dim=1) # 扩展注意力掩码 prefix_attention_mask = torch.ones( (batch_size, self.num_prefix), dtype=attention_mask.dtype, device=device ) extended_attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) # ✨ 关键修复: 处理标签 if labels is not None: # 为前缀部分创建-100的标签（表示忽略） prefix_labels = torch.full( (batch_size, self.num_prefix), fill_value=-100, # -100表示忽略这些位置的损失 dtype=labels.dtype, device=device ) # 扩展标签 extended_labels = torch.cat((prefix_labels, labels), dim=1) else: extended_labels = None # 调试输出 # print(f"[DEBUG] 原始输入大小: {input_ids.shape}") # print(f"[DEBUG] 扩展嵌入大小: {input_embeds.shape}") # print(f"[DEBUG] 扩展掩码大小: {extended_attention_mask.shape}") # if extended_labels is not None: # print(f"[DEBUG] 扩展标签大小: {extended_labels.shape}") # 确保不提供input_ids if 'input_ids' in kwargs: del kwargs['input_ids'] # ✨ 传递扩展后的标签 return self.model( inputs_embeds=input_embeds, attention_mask=extended_attention_mask, labels=extended_labels, # 这是关键修改 use_cache=False, **kwargs) # class PrefixKGEmbedding(nn.Module): # def __init__( # self, # num_ent, # num_rel, # dim_llm, # num_prefix # ): # super(PrefixKGEmbedding, self).__init__() # self.emb_dim = num_prefix * dim_llm # self.ent_embeddings = nn.Embedding(num_ent, self.emb_dim) # self.rel_embeddings = nn.Embedding(num_rel, self.emb_dim) # # # def forward(self, triple_ids): # head, relation, tail = triple_ids[:, 0], triple_ids[:, 1], triple_ids[:, 2] # h = self.ent_embeddings(head) # r = self.rel_embeddings(relation) # t = self.ent_embeddings(tail) # prefix = torch.stack((h, r, t), dim=1) # return prefix class PretrainKGEmbedding(nn.Module): def __init__( self, pretrain_ent_embs, pretrain_rel_embs, dim_llm, num_prefix ): super(PretrainKGEmbedding, self).__init__() self.num_prefix = num_prefix self.llm_dim = dim_llm self.emb_dim = num_prefix * dim_llm self.ent_embeddings = nn.Embedding.from_pretrained(pretrain_ent_embs) self.rel_embeddings = nn.Embedding.from_pretrained(pretrain_rel_embs) self.pretrain_dim = self.ent_embeddings.weight.shape[1] # Froze the pretrain embeddings self.ent_embeddings.requires_grad_(False) self.rel_embeddings.requires_grad_(False) self.adapter = nn.Linear(self.pretrain_dim, self.emb_dim) def forward(self, triple_ids): # main training stage if triple_ids.shape[1] == 3: head, relation, tail = triple_ids[:, 0], triple_ids[:, 1], triple_ids[:, 2] h = self.ent_embeddings(head) r = self.rel_embeddings(relation) t = self.ent_embeddings(tail) pretrain_embs = torch.stack((h, r, t), dim=1) prefix = self.adapter(pretrain_embs).reshape(-1, 3*self.num_prefix, self.llm_dim) return prefix # entity-aware pre-funing else: ent = triple_ids.reshape(-1,) emb = self.ent_embeddings(ent) prefix = self.adapter(emb).reshape(-1, self.num_prefix, self.llm_dim) # print(prefix.shape) return prefix