如何仅用源代码通过CodeBERT生成嵌入？现有方法效果存疑

问题说明 #

CodeBERT官方示例都是用自然语言-代码对生成嵌入，但我现在需要只靠无关联自然语言的源代码生成嵌入，用来计算代码间的余弦相似度。之前试过移除自然语言令牌的方式生成嵌入，但相似度结果很差，不确定是不是最优方案。另外发现code-bert-score支持把自然语言设为可选输入，但搞不懂它具体是怎么实现的。

当前代码的问题 #

我之前构造的令牌序列是[<s>] + [</s>] + 代码令牌 + [</s>]，这种结构不符合CodeBERT预训练时的单代码输入格式。CodeBERT预训练时，单代码输入的标准格式应该是[<s>] + 代码令牌 + [</s>]，多出来的</s>会干扰模型提取代码特征，导致嵌入质量下降，自然相似度结果就不好。

优化后的代码实现 #

按照预训练的标准格式构造输入，直接用<s>开头，代码令牌跟上，最后用</s>结尾，同时用torch.no_grad()关闭梯度计算节省资源。另外取<s>令牌的嵌入作为代码的整体表示，这是CodeBERT常用的做法：

from transformers import AutoTokenizer, AutoModel
import torch
from sklearn.metrics.pairwise import cosine_similarity

# 加载CodeBERT模型和分词器
tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
model = AutoModel.from_pretrained("microsoft/codebert-base")

def get_code_embedding(code):
    # 对代码进行分词
    code_tokens = tokenizer.tokenize(code)
    # 构造标准输入序列：<s> + 代码令牌 + </s>
    tokens = [tokenizer.cls_token] + code_tokens + [tokenizer.sep_token]
    # 转换为token ID并添加batch维度
    token_ids = tokenizer.convert_tokens_to_ids(tokens)
    input_ids = torch.tensor([token_ids])
    
    # 获取模型输出，取第一个令牌（<s>）的向量作为代码嵌入
    with torch.no_grad():
        outputs = model(input_ids)
        cls_embedding = outputs[0][:, 0, :].numpy()
    
    return cls_embedding

# 测试两个相似代码的相似度
code1 = "def max(a,b): if a>b: return a else return b"
code2 = "def maximum(x,y): return x if x>y else y"

emb1 = get_code_embedding(code1)
emb2 = get_code_embedding(code2)

similarity = cosine_similarity(emb1, emb2)[0][0]
print(f"代码相似度：{similarity:.4f}")

code-bert-score的单代码输入逻辑 #

code-bert-score里处理可选自然语言输入的逻辑很简单：当只给代码时，就构造单模态的代码输入序列，不用双模态的NL-PL对格式。核心就是遵循CodeBERT预训练的单代码输入结构，要么取<s>令牌的嵌入，要么对所有代码令牌的嵌入做均值池化，得到代码的整体表示。

额外优化技巧 #

• 均值池化替代CLS令牌：有时候对代码部分的令牌嵌入做均值池化，能得到更稳定的代码表示，适合长代码场景：

  def get_code_embedding_avg(code):
      code_tokens = tokenizer.tokenize(code)
      tokens = [tokenizer.cls_token] + code_tokens + [tokenizer.sep_token]
      token_ids = tokenizer.convert_tokens_to_ids(tokens)
      input_ids = torch.tensor([token_ids])
      
      with torch.no_grad():
          outputs = model(input_ids)
          # 提取代码部分的嵌入（排除首尾的特殊令牌），计算均值
          code_part_embeddings = outputs[0][:, 1:-1, :]
          avg_embedding = torch.mean(code_part_embeddings, dim=1).numpy()
      
      return avg_embedding
  

• 统一输入长度：对于较长的代码，用分词器的padding和truncation参数把输入统一到固定长度（比如512），避免模型处理不同长度序列时的性能波动：

  inputs = tokenizer(code, return_tensors="pt", padding="max_length", truncation=True, max_length=512)
  with torch.no_grad():
      outputs = model(**inputs)
  

内容的提问来源于stack exchange，提问作者Armand Mousavi