Upload folder using huggingface_hub

README.md

@@ -1,3 +1,13 @@
+---
+language: ko
+tags:
+- text-classification
+- steel-industry
+- tf-idf
+- neural-network
+license: mit
+---
+
 # Steel Industry Material Classification Model
 
 This model is trained to classify steel industry materials and products based on text descriptions. It uses a custom TF-IDF + Neural Network approach and can classify input text into 66 different steel-related categories.
@@ -25,37 +35,70 @@ The model can classify the following steel industry materials:
 
 ```python
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
-import pickle
 import joblib
-from sklearn.feature_extraction.text import TfidfVectorizer
-
-# Load model components
-with open('vectorizer.pkl', 'rb') as f:
-    vectorizer = joblib.load(f)
-
-with open('model.pkl', 'rb') as f:
-    model_data = pickle.load(f)
-
-model = model_data['model']
-id2label = model_data['id2label']
-
-# Prepare input
+import json
+from huggingface_hub import hf_hub_download
+
+# SimpleClassifier 클래스 정의
+class SimpleClassifier(nn.Module):
+    def __init__(self, input_size, num_classes):
+        super(SimpleClassifier, self).__init__()
+        self.fc1 = nn.Linear(input_size, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, num_classes)
+        self.dropout = nn.Dropout(0.3)
+        
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = F.relu(self.fc2(x))
+        x = self.dropout(x)
+        x = self.fc3(x)
+        return x
+
+# 모델 파일들 다운로드
+model_path = hf_hub_download(repo_id="Halfotter/flud", filename="pytorch_model.bin")
+vectorizer_path = hf_hub_download(repo_id="Halfotter/flud", filename="vectorizer.pkl")
+config_path = hf_hub_download(repo_id="Halfotter/flud", filename="config.json")
+
+# 설정 로드
+with open(config_path, 'r', encoding='utf-8') as f:
+    config = json.load(f)
+
+id2label = config.get('id2label', {})
+
+# 모델 로드
+input_size = 3000  # TF-IDF 특성 수
+num_classes = len(id2label)
+model = SimpleClassifier(input_size, num_classes)
+model.load_state_dict(torch.load(model_path, map_location='cpu'))
+
+# 벡터라이저 로드
+vectorizer = joblib.load(vectorizer_path)
+
+# 예측 함수
+def predict(text):
+    model.eval()
+    
+    # TF-IDF 벡터화
+    text_vector = vectorizer.transform([text]).toarray()
+    text_tensor = torch.FloatTensor(text_vector)
+    
+    with torch.no_grad():
+        outputs = model(text_tensor)
+        probabilities = F.softmax(outputs, dim=1)
+        predicted_class = torch.argmax(probabilities, dim=1).item()
+    
+    label = id2label[str(predicted_class)]
+    confidence = probabilities[0][predicted_class].item()
+    
+    return label, confidence
+
+# 사용 예시
 text = "철광석을 고로에서 환원하여 선철을 제조하는 과정"
-text_vector = vectorizer.transform([text]).toarray()
-text_tensor = torch.FloatTensor(text_vector)
-
-# Predict
-model.eval()
-with torch.no_grad():
-    outputs = model(text_tensor)
-    probabilities = F.softmax(outputs, dim=1)
-    predicted_class = torch.argmax(probabilities, dim=1).item()
-
-# Get label
-label = id2label[str(predicted_class)]
-confidence = probabilities[0][predicted_class].item()
-
+label, confidence = predict(text)
 print(f"Predicted: {label}")
 print(f"Confidence: {confidence:.4f}")
 ```

requirements.txt

@@ -1,5 +1,6 @@
-torch>=1.9.0
-scikit-learn>=1.0.0
-numpy>=1.21.0
-pandas>=1.3.0
-joblib>=1.1.0
+torch==2.8.0
+scikit-learn==1.7.1
+joblib==1.5.2
+numpy==2.2.6
+pandas==2.3.2
+huggingface-hub>=0.19.0

test_current_model.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pickle
+import joblib
+import numpy as np
+
+# SimpleClassifier 클래스 정의
+class SimpleClassifier(nn.Module):
+    def __init__(self, input_size, num_classes):
+        super(SimpleClassifier, self).__init__()
+        self.fc1 = nn.Linear(input_size, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, num_classes)
+        self.dropout = nn.Dropout(0.3)
+        
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = F.relu(self.fc2(x))
+        x = self.dropout(x)
+        x = self.fc3(x)
+        return x
+
+def test_current_model():
+    """현재 모델 테스트"""
+    print("=== 현재 모델 테스트 ===")
+    
+    try:
+        # 설정 로드
+        with open('config.json', 'r', encoding='utf-8') as f:
+            import json
+            config = json.load(f)
+        
+        id2label = config.get('id2label', {})
+        print(f"라벨 수: {len(id2label)}")
+        
+        # 모델 로드
+        input_size = 3000  # TF-IDF 특성 수
+        num_classes = len(id2label)
+        model = SimpleClassifier(input_size, num_classes)
+        model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
+        
+        # 벡터라이저 로드
+        vectorizer = joblib.load('vectorizer.pkl')
+        
+        model.eval()
+        
+        # 테스트 단어들 (환원철 포함)
+        test_words = ["철ㄹ", "CaO", "해면철", "등류", "환원철"]
+        
+        for word in test_words:
+            print(f"\n{'='*50}")
+            print(f"입력: '{word}'")
+            print(f"{'='*50}")
+            
+            # TF-IDF 벡터화
+            word_vector = vectorizer.transform([word]).toarray()
+            word_tensor = torch.FloatTensor(word_vector)
+            
+            with torch.no_grad():
+                outputs = model(word_tensor)
+                probabilities = F.softmax(outputs, dim=1)
+                
+                # 상위 5개 예측
+                top_probs, top_indices = torch.topk(probabilities, 5, dim=1)
+                
+                print(f"최대 확률: {probabilities.max().item():.4f} ({probabilities.max().item()*100:.1f}%)")
+                print(f"상위 5개 예측:")
+                
+                for i in range(5):
+                    label_id = top_indices[0][i].item()
+                    probability = top_probs[0][i].item()
+                    label = id2label.get(str(label_id), f"Unknown_{label_id}")
+                    print(f"  {i+1}. {label}: {probability:.4f} ({probability*100:.1f}%)")
+                    
+    except Exception as e:
+        print(f"에러 발생: {e}")
+        import traceback
+        traceback.print_exc()
+
+if __name__ == "__main__":
+    test_current_model()

test_sintered_ore.py

@@ -0,0 +1,78 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pickle
+import joblib
+import numpy as np
+
+# SimpleClassifier 클래스 정의
+class SimpleClassifier(nn.Module):
+    def __init__(self, input_size, num_classes):
+        super(SimpleClassifier, self).__init__()
+        self.fc1 = nn.Linear(input_size, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, num_classes)
+        self.dropout = nn.Dropout(0.3)
+        
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = F.relu(self.fc2(x))
+        x = self.dropout(x)
+        x = self.fc3(x)
+        return x
+
+def test_sintered_ore():
+    """소결광 테스트"""
+    print("=== 소결광 테스트 ===")
+    
+    try:
+        # 설정 로드
+        with open('config.json', 'r', encoding='utf-8') as f:
+            import json
+            config = json.load(f)
+        
+        id2label = config.get('id2label', {})
+        
+        # 모델 로드
+        input_size = 3000
+        num_classes = len(id2label)
+        model = SimpleClassifier(input_size, num_classes)
+        model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
+        
+        # 벡터라이저 로드
+        vectorizer = joblib.load('vectorizer.pkl')
+        
+        model.eval()
+        
+        # 소결광 테스트
+        test_word = "소결광"
+        print(f"입력: '{test_word}'")
+        
+        # TF-IDF 벡터화
+        word_vector = vectorizer.transform([test_word]).toarray()
+        word_tensor = torch.FloatTensor(word_vector)
+        
+        with torch.no_grad():
+            outputs = model(word_tensor)
+            probabilities = F.softmax(outputs, dim=1)
+            
+            # 상위 10개 예측
+            top_probs, top_indices = torch.topk(probabilities, 10, dim=1)
+            
+            print(f"최대 확률: {probabilities.max().item():.4f} ({probabilities.max().item()*100:.1f}%)")
+            print(f"상위 10개 예측:")
+            
+            for i in range(10):
+                label_id = top_indices[0][i].item()
+                probability = top_probs[0][i].item()
+                label = id2label.get(str(label_id), f"Unknown_{label_id}")
+                print(f"  {i+1}. {label}: {probability:.4f} ({probability*100:.1f}%)")
+                
+    except Exception as e:
+        print(f"에러 발생: {e}")
+        import traceback
+        traceback.print_exc()
+
+if __name__ == "__main__":
+    test_sintered_ore()