Halfotter/flud

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import numpy as np
5import json
6import os
7from transformers import PreTrainedModel, PretrainedConfig, XLMRobertaModel, XLMRobertaConfig
8
9class XLMSteelConfig(PretrainedConfig):
10    """XLM-RoBERTa 철강 분류기 설정"""
11    model_type = "xlm_steel_classifier"
12    
13    def __init__(self, num_labels=66, **kwargs):
14        super().__init__(**kwargs)
15        self.num_labels = num_labels
16
17class XLMIntegratedModel(PreTrainedModel):
18    """XLM-RoBERTa + TF-IDF 통합 모델"""
19    config_class = XLMSteelConfig
20    
21    def __init__(self, config):
22        super().__init__(config)
23        
24        # XLM-RoBERTa 모델
25        self.xlm_roberta = XLMRobertaModel.from_pretrained('xlm-roberta-base')
26        
27        # TF-IDF 벡터라이저 정보 저장
28        self.feature_names = getattr(config, 'feature_names', [])
29        self.input_size = getattr(config, 'input_size', 3000)
30        
31        # 신경망 레이어 (기존 TF-IDF 모델 구조)
32        self.fc1 = nn.Linear(self.input_size, 256)
33        self.fc2 = nn.Linear(256, 128)
34        self.fc3 = nn.Linear(128, config.num_labels)
35        self.dropout = nn.Dropout(0.3)
36        
37        # 라벨 매핑 저장
38        self.id2label = config.id2label
39        self.num_classes = config.num_labels
40        
41        # 벡터라이저의 특성 정보를 텐서로 저장
42        self.register_buffer('feature_names_list', torch.tensor([hash(f) for f in self.feature_names], dtype=torch.long))
43        
44    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
45        """통합 forward"""
46        # XLM-RoBERTa 출력
47        if input_ids is not None:
48            xlm_outputs = self.xlm_roberta(
49                input_ids=input_ids,
50                attention_mask=attention_mask,
51                return_dict=True
52            )
53            xlm_features = xlm_outputs.pooler_output
54        else:
55            xlm_features = torch.zeros(1, self.xlm_roberta.config.hidden_size)
56        
57        # TF-IDF 벡터화 (내부적으로 수행)
58        if input_ids is not None:
59            # input_ids를 텍스트로 변환하여 TF-IDF 벡터화
60            text_vector = self._vectorize_from_ids(input_ids[0])
61            tfidf_features = torch.FloatTensor(text_vector).unsqueeze(0)
62        else:
63            tfidf_features = torch.zeros(1, self.input_size)
64        
65        # 신경망 통과 (TF-IDF 부분만 사용)
66        x = F.relu(self.fc1(tfidf_features))
67        x = self.dropout(x)
68        x = F.relu(self.fc2(x))
69        x = self.dropout(x)
70        logits = self.fc3(x)
71        
72        # 손실 계산
73        loss = None
74        if labels is not None:
75            loss_fct = nn.CrossEntropyLoss()
76            loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
77        
78        return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}
79    
80    def _vectorize_from_ids(self, input_ids):
81        """input_ids를 TF-IDF 벡터로 변환"""
82        vector = np.zeros(self.input_size)
83        
84        # input_ids를 기반으로 벡터 생성
85        for token_id in input_ids:
86            if token_id < self.input_size:
87                vector[token_id] += 1
88        
89        if np.sum(vector) > 0:
90            vector = vector / np.sum(vector)
91        
92        return vector
93
94# 전역 변수
95model = None
96
97def load_model():
98    """모델 로드"""
99    global model
100    
101    # 설정 파일 로드
102    config_path = os.path.join(os.getcwd(), "config.json")
103    with open(config_path, 'r', encoding='utf-8') as f:
104        config_data = json.load(f)
105    
106    # XLMSteelConfig 생성
107    config = XLMSteelConfig(
108        num_labels=config_data['num_labels'],
109        id2label=config_data['id2label'],
110        label2id=config_data['label2id'],
111        feature_names=config_data.get('feature_names', []),
112        input_size=config_data.get('input_size', 3000)
113    )
114    
115    # 모델 생성 및 로드
116    model = XLMIntegratedModel(config)
117    model_path = os.path.join(os.getcwd(), "xlm_integrated_model.bin")
118    model.load_state_dict(torch.load(model_path, map_location='cpu'))
119    model.eval()
120    
121    return model
122
123def predict(inputs):
124    """예측 함수"""
125    global model
126    if model is None:
127        model = load_model()
128    
129    # 입력 처리
130    if isinstance(inputs, str):
131        text = inputs
132    elif isinstance(inputs, list):
133        text = inputs[0] if len(inputs) > 0 else ""
134    elif isinstance(inputs, dict) and "inputs" in inputs:
135        text = inputs["inputs"]
136    else:
137        text = str(inputs)
138    
139    # 텍스트를 토큰 ID로 변환 (간단한 구현)
140    tokens = text.lower().split()
141    input_ids = torch.tensor([[hash(token) % 50000 for token in tokens]])  # XLM-RoBERTa vocab size
142    attention_mask = torch.ones_like(input_ids)
143    
144    # 예측
145    with torch.no_grad():
146        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
147        logits = outputs["logits"]
148        probabilities = F.softmax(logits, dim=1)
149        predicted_class = torch.argmax(probabilities, dim=1).item()
150        
151        label = model.id2label[str(predicted_class)]
152        confidence = probabilities[0][predicted_class].item()
153        
154        return {
155            "label": label,
156            "confidence": confidence,
157            "text": text
158        }
159
160# 모델 초기 로드
161if __name__ == "__main__":
162    load_model()
163