Halfotter/flud

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import pickle
5import joblib
6import numpy as np
7
8# SimpleClassifier 클래스 정의
9class SimpleClassifier(nn.Module):
10    def __init__(self, input_size, num_classes):
11        super(SimpleClassifier, self).__init__()
12        self.fc1 = nn.Linear(input_size, 256)
13        self.fc2 = nn.Linear(256, 128)
14        self.fc3 = nn.Linear(128, num_classes)
15        self.dropout = nn.Dropout(0.3)
16        
17    def forward(self, x):
18        x = F.relu(self.fc1(x))
19        x = self.dropout(x)
20        x = F.relu(self.fc2(x))
21        x = self.dropout(x)
22        x = self.fc3(x)
23        return x
24
25def test_current_model():
26    """현재 모델 테스트"""
27    print("=== 현재 모델 테스트 ===")
28    
29    try:
30        # 설정 로드
31        with open('config.json', 'r', encoding='utf-8') as f:
32            import json
33            config = json.load(f)
34        
35        id2label = config.get('id2label', {})
36        print(f"라벨 수: {len(id2label)}")
37        
38        # 모델 로드
39        input_size = 3000  # TF-IDF 특성 수
40        num_classes = len(id2label)
41        model = SimpleClassifier(input_size, num_classes)
42        model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
43        
44        # 벡터라이저 로드
45        vectorizer = joblib.load('vectorizer.pkl')
46        
47        model.eval()
48        
49        # 테스트 단어들 (환원철 포함)
50        test_words = ["철ㄹ", "CaO", "해면철", "등류", "환원철"]
51        
52        for word in test_words:
53            print(f"\n{'='*50}")
54            print(f"입력: '{word}'")
55            print(f"{'='*50}")
56            
57            # TF-IDF 벡터화
58            word_vector = vectorizer.transform([word]).toarray()
59            word_tensor = torch.FloatTensor(word_vector)
60            
61            with torch.no_grad():
62                outputs = model(word_tensor)
63                probabilities = F.softmax(outputs, dim=1)
64                
65                # 상위 5개 예측
66                top_probs, top_indices = torch.topk(probabilities, 5, dim=1)
67                
68                print(f"최대 확률: {probabilities.max().item():.4f} ({probabilities.max().item()*100:.1f}%)")
69                print(f"상위 5개 예측:")
70                
71                for i in range(5):
72                    label_id = top_indices[0][i].item()
73                    probability = top_probs[0][i].item()
74                    label = id2label.get(str(label_id), f"Unknown_{label_id}")
75                    print(f"  {i+1}. {label}: {probability:.4f} ({probability*100:.1f}%)")
76                    
77    except Exception as e:
78        print(f"에러 발생: {e}")
79        import traceback
80        traceback.print_exc()
81
82if __name__ == "__main__":
83    test_current_model()
84