FastestAI/CodeBert_Redundant_Detection_Task

1import torch2import torch.nn as nn3from transformers import AutoModel4import re5 6class CodeSimilarityClassifier(nn.Module):7    def __init__(self, model_name="microsoft/codebert-base", num_labels=3):8        super().__init__()9        self.encoder = AutoModel.from_pretrained(model_name)10        self.dropout = nn.Dropout(0.1)11 12        # Create a more powerful classification head13        hidden_size = self.encoder.config.hidden_size14 15        self.classifier = nn.Sequential(16            nn.Linear(hidden_size, hidden_size),17            nn.LayerNorm(hidden_size),18            nn.GELU(),19            nn.Dropout(0.1),20            nn.Linear(hidden_size, 512),21            nn.LayerNorm(512),22            nn.GELU(),23            nn.Dropout(0.1),24            nn.Linear(512, num_labels)25        )26 27    def forward(self, input_ids, attention_mask):28        outputs = self.encoder(29            input_ids=input_ids,30            attention_mask=attention_mask,31            return_dict=True32        )33 34        pooled_output = outputs.pooler_output35        logits = self.classifier(pooled_output)36 37        return logits38 39def extract_features(source_code, test_code_1, test_code_2):40    """Extract specific features to help the model identify similarities"""41 42    # Extract test fixtures43    fixture1 = re.search(r'TEST(?:_F)?\s*\(\s*(\w+)', test_code_1)44    fixture1 = fixture1.group(1) if fixture1 else ""45 46    fixture2 = re.search(r'TEST(?:_F)?\s*\(\s*(\w+)', test_code_2)47    fixture2 = fixture2.group(1) if fixture2 else ""48 49    # Extract test names50    name1 = re.search(r'TEST(?:_F)?\s*\(\s*\w+\s*,\s*(\w+)', test_code_1)51    name1 = name1.group(1) if name1 else ""52 53    name2 = re.search(r'TEST(?:_F)?\s*\(\s*\w+\s*,\s*(\w+)', test_code_2)54    name2 = name2.group(1) if name2 else ""55 56    # Extract assertions57    assertions1 = re.findall(r'(EXPECT_|ASSERT_)(\w+)', test_code_1)58    assertions2 = re.findall(r'(EXPECT_|ASSERT_)(\w+)', test_code_2)59 60    # Extract function/method calls61    calls1 = re.findall(r'(\w+)\s*\(', test_code_1)62    calls2 = re.findall(r'(\w+)\s*\(', test_code_2)63 64    # Create explicit feature section65    same_fixture = "SAME_FIXTURE" if fixture1 == fixture2 else "DIFFERENT_FIXTURE"66    common_assertions = set([a[0] + a[1] for a in assertions1]).intersection(set([a[0] + a[1] for a in assertions2]))67    common_calls = set(calls1).intersection(set(calls2))68 69    features = (70        f"METADATA: {same_fixture} | "71        f"FIXTURE1: {fixture1} | FIXTURE2: {fixture2} | "72        f"NAME1: {name1} | NAME2: {name2} | "73        f"COMMON_ASSERTIONS: {len(common_assertions)} | "74        f"COMMON_CALLS: {len(common_calls)} | "75        f"ASSERTION_RATIO: {len(common_assertions)/(len(assertions1) + len(assertions2)) if assertions1 and assertions2 else 0}"76    )77 78    return features