Upload inference.py with huggingface_hub

inference.py

@@ -4,54 +4,66 @@ import torch.nn.functional as F
 import numpy as np
 import json
 import os
-from transformers import PreTrainedModel, PretrainedConfig
+from transformers import PreTrainedModel, PretrainedConfig, XLMRobertaModel, XLMRobertaConfig
 
-class SteelConfig(PretrainedConfig):
-    """철강 분류기 설정"""
-    model_type = "steel_classifier"
+class XLMSteelConfig(PretrainedConfig):
+    """XLM-RoBERTa 철강 분류기 설정"""
+    model_type = "xlm_steel_classifier"
     
-    def __init__(self, input_size=3000, hidden_size=256, intermediate_size=128, num_labels=66, **kwargs):
+    def __init__(self, num_labels=66, **kwargs):
         super().__init__(**kwargs)
-        self.input_size = input_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
         self.num_labels = num_labels
 
-class SteelModel(PreTrainedModel):
-    """Hugging Face 표준 철강 분류기"""
-    config_class = SteelConfig
+class XLMIntegratedModel(PreTrainedModel):
+    """XLM-RoBERTa + TF-IDF 통합 모델"""
+    config_class = XLMSteelConfig
     
     def __init__(self, config):
         super().__init__(config)
         
-        # 신경망 레이어
-        self.fc1 = nn.Linear(config.input_size, config.hidden_size)
-        self.fc2 = nn.Linear(config.hidden_size, config.intermediate_size)
-        self.fc3 = nn.Linear(config.intermediate_size, config.num_labels)
+        # XLM-RoBERTa 모델
+        self.xlm_roberta = XLMRobertaModel.from_pretrained('xlm-roberta-base')
+        
+        # TF-IDF 벡터라이저 정보 저장
+        self.feature_names = getattr(config, 'feature_names', [])
+        self.input_size = getattr(config, 'input_size', 3000)
+        
+        # 신경망 레이어 (기존 TF-IDF 모델 구조)
+        self.fc1 = nn.Linear(self.input_size, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, config.num_labels)
         self.dropout = nn.Dropout(0.3)
         
-        # 라벨 매핑
-        self.id2label = config.id2label if hasattr(config, 'id2label') else {}
-        self.label2id = config.label2id if hasattr(config, 'label2id') else {}
+        # 라벨 매핑 저장
+        self.id2label = config.id2label
+        self.num_classes = config.num_labels
+        
+        # 벡터라이저의 특성 정보를 텐서로 저장
+        self.register_buffer('feature_names_list', torch.tensor([hash(f) for f in self.feature_names], dtype=torch.long))
         
     def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
-        """표준 Hugging Face forward"""
-        # TF-IDF 벡터화 (input_ids를 텍스트로 변환)
+        """통합 forward"""
+        # XLM-RoBERTa 출력
         if input_ids is not None:
-            # input_ids를 텍스트로 변환하는 로직
-            batch_size = input_ids.shape[0]
-            features = torch.zeros(batch_size, self.config.input_size)
-            
-            for i in range(batch_size):
-                # 간단한 벡터화 (실제로는 더 정교한 TF-IDF 필요)
-                text_vector = self._simple_vectorize(input_ids[i])
-                features[i] = torch.FloatTensor(text_vector)
+            xlm_outputs = self.xlm_roberta(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                return_dict=True
+            )
+            xlm_features = xlm_outputs.pooler_output
         else:
-            # 기본값
-            features = torch.zeros(1, self.config.input_size)
+            xlm_features = torch.zeros(1, self.xlm_roberta.config.hidden_size)
         
-        # 신경망 통과
-        x = F.relu(self.fc1(features))
+        # TF-IDF 벡터화 (내부적으로 수행)
+        if input_ids is not None:
+            # input_ids를 텍스트로 변환하여 TF-IDF 벡터화
+            text_vector = self._vectorize_from_ids(input_ids[0])
+            tfidf_features = torch.FloatTensor(text_vector).unsqueeze(0)
+        else:
+            tfidf_features = torch.zeros(1, self.input_size)
+        
+        # 신경망 통과 (TF-IDF 부분만 사용)
+        x = F.relu(self.fc1(tfidf_features))
         x = self.dropout(x)
         x = F.relu(self.fc2(x))
         x = self.dropout(x)
@@ -65,12 +77,13 @@ class SteelModel(PreTrainedModel):
         
         return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}
     
-    def _simple_vectorize(self, input_ids):
-        """간단한 벡터화"""
-        vector = np.zeros(self.config.input_size)
+    def _vectorize_from_ids(self, input_ids):
+        """input_ids를 TF-IDF 벡터로 변환"""
+        vector = np.zeros(self.input_size)
+        
         # input_ids를 기반으로 벡터 생성
         for token_id in input_ids:
-            if token_id < self.config.input_size:
+            if token_id < self.input_size:
                 vector[token_id] += 1
         
         if np.sum(vector) > 0:
@@ -90,19 +103,18 @@ def load_model():
     with open(config_path, 'r', encoding='utf-8') as f:
         config_data = json.load(f)
     
-    # SteelConfig 생성
-    config = SteelConfig(
-        input_size=config_data['input_size'],
-        hidden_size=config_data['hidden_size'],
-        intermediate_size=config_data['intermediate_size'],
+    # XLMSteelConfig 생성
+    config = XLMSteelConfig(
         num_labels=config_data['num_labels'],
         id2label=config_data['id2label'],
-        label2id=config_data['label2id']
+        label2id=config_data['label2id'],
+        feature_names=config_data.get('feature_names', []),
+        input_size=config_data.get('input_size', 3000)
     )
     
     # 모델 생성 및 로드
-    model = SteelModel(config)
-    model_path = os.path.join(os.getcwd(), "integrated_model.bin")
+    model = XLMIntegratedModel(config)
+    model_path = os.path.join(os.getcwd(), "xlm_integrated_model.bin")
     model.load_state_dict(torch.load(model_path, map_location='cpu'))
     model.eval()
     
@@ -124,13 +136,14 @@ def predict(inputs):
     else:
         text = str(inputs)
     
-    # 텍스트를 간단한 토큰 ID로 변환
+    # 텍스트를 토큰 ID로 변환 (간단한 구현)
     tokens = text.lower().split()
-    input_ids = torch.tensor([[hash(token) % model.config.input_size for token in tokens]])
+    input_ids = torch.tensor([[hash(token) % 50000 for token in tokens]])  # XLM-RoBERTa vocab size
+    attention_mask = torch.ones_like(input_ids)
     
     # 예측
     with torch.no_grad():
-        outputs = model(input_ids=input_ids)
+        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
         logits = outputs["logits"]
         probabilities = F.softmax(logits, dim=1)
         predicted_class = torch.argmax(probabilities, dim=1).item()