initial push (safetensors + custom head)

README.md

@@ -1,54 +1,6 @@
----
-language: 
-- code
-library_name: transformers
-pipeline_tag: text-classification
-tags:
-- code-analysis
-- vulnerability-detection
-- security  
-- cwe
-license: mit
-base_model: microsoft/codebert-base
----
-
-# CodeBERT Vulnerability Detector (Multi-class)
-
-C/C++ 코드의 취약점을 탐지하는 다중 클래스 분류 모델입니다.
-
-## 모델 정보
-- **기반 모델**: microsoft/codebert-base
-- **분류 클래스**: 4개 (CWE-79, CWE-89, CWE-119, 기타)
-- **입력**: C/C++ 소스 코드 텍스트
-
-## 사용 방법
-
+# Custom RoBERTa (safetensors)
+Load with:
 ```python
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-
-# 모델 로드
-model_name = "eunJ/codebert_vulnerability_detector_multi"
-tokenizer = AutoTokenizer.from_pretrained(model_name)  
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-
-# 코드 분석
-code = '''
-char buffer[100];
-gets(buffer);  
-'''
-
-inputs = tokenizer(code, return_tensors="pt", max_length=512, truncation=True)
-with torch.no_grad():
-    outputs = model(**inputs)
-    predictions = torch.softmax(outputs.logits, dim=-1)
-    predicted_class = torch.argmax(predictions)
-
-print(f"예측 클래스: {predicted_class.item()}")
-```
-
-## 클래스 레이블
-- 0: CWE-79 (Cross-site Scripting)
-- 1: CWE-89 (SQL Injection) 
-- 2: CWE-119 (Buffer Overflow)
-- 3: CWE-Other (기타)
+tok = AutoTokenizer.from_pretrained("eunJ/codebert_vulnerability_detector_multi")
+model = AutoModelForSequenceClassification.from_pretrained("eunJ/codebert_vulnerability_detector_multi", trust_remote_code=True).eval()

config.json

@@ -1,19 +1,42 @@
 {
-  "model_type": "roberta",
+  "_name_or_path": "microsoft/codebert-base",
   "architectures": [
-    "RobertaForSequenceClassification"
+    "RobertaModel"
   ],
-  "num_labels": 4,
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoModelForSequenceClassification": "modeling_my_roberta.MyRobertaForSequenceClassification"
+  },
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
   "id2label": {
     "0": "LABEL_0",
     "1": "LABEL_1",
     "2": "LABEL_2",
     "3": "LABEL_3"
   },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
   "label2id": {
     "LABEL_0": 0,
     "LABEL_1": 1,
     "LABEL_2": 2,
     "LABEL_3": 3
-  }
-}
+  },
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 514,
+  "model_type": "roberta",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "output_past": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.49.0",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 50265
+}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c8037175a0f7980967910a796119cf0250a6c0200f7568225f2e9aaeb43b9b68
-size 498633008
+oid sha256:28416b8a21440d53974e0456bf036dd7b9778a371c50b159f8870f6742c8ecad
+size 496256360

modeling_my_roberta.py

@@ -0,0 +1,54 @@
+# modeling_my_roberta.py
+from typing import Optional, Tuple
+import torch
+import torch.nn as nn
+from transformers import RobertaModel, RobertaPreTrainedModel
+
+class MyRobertaForSequenceClassification(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = getattr(config, "num_labels", 4)
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.classifier = nn.Linear(config.hidden_size, self.num_labels)
+        self.loss_fn = nn.CrossEntropyLoss()
+        self.post_init()  # init_weights
+
+    def _get_pad_id(self, input_ids):
+        # Roberta pad = 1가 일반적이지만 토크나이저에 맞춰 attention_mask 직접 받는 게 더 안전
+        return 1
+
+    def _pool(self, last_hidden_state, attention_mask, model_type: str = "roberta"):
+        # 너의 로직: roberta면 CLS, 그 외 mask mean
+        if last_hidden_state.dim() == 3:
+            if model_type in {"bert", "roberta", "deberta", "xlm-roberta", "electra"}:
+                return last_hidden_state[:, 0, :]
+            mask = attention_mask.unsqueeze(-1).float()
+            summed = (last_hidden_state * mask).sum(dim=1)
+            denom = mask.sum(dim=1).clamp(min=1e-9)
+            return summed / denom
+        elif last_hidden_state.dim() == 2:
+            return last_hidden_state
+        else:
+            raise ValueError(f"Unexpected hidden dim: {last_hidden_state.dim()}")
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs
+    ):
+        if attention_mask is None and input_ids is not None:
+            pad_id = self._get_pad_id(input_ids)
+            attention_mask = input_ids.ne(pad_id)
+
+        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
+        hidden = outputs.last_hidden_state
+        pooled = self._pool(hidden, attention_mask, "roberta")
+        logits = self.classifier(pooled)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_fn(logits, labels.long())
+
+        return {"loss": loss, "logits": logits}

tokenizer_config.json

@@ -45,7 +45,6 @@
   "bos_token": "<s>",
   "clean_up_tokenization_spaces": false,
   "cls_token": "<s>",
-  "do_lower_case": false,
   "eos_token": "</s>",
   "errors": "replace",
   "extra_special_tokens": {},
@@ -54,5 +53,6 @@
   "pad_token": "<pad>",
   "sep_token": "</s>",
   "tokenizer_class": "RobertaTokenizer",
+  "trim_offsets": true,
   "unk_token": "<unk>"
 }