Add Classification

app.py

@@ -13,7 +13,7 @@ os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 os.environ["TF_NUM_INTEROP_THREADS"] = "1"
 os.environ["TF_NUM_INTRAOP_THREADS"] = "1"
 
-app = FastAPI(title="NER API Service")
+app = FastAPI(title="NER & Classification API Service")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -22,6 +22,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Global Variables NER
 ner_model = None
 vocab = None
 MAX_SEQ_LEN = 64
@@ -30,39 +31,32 @@ UNK_TOKEN = "<UNK>"
 NUM_TOKEN = "<NUM>"
 PRICE_TOKEN = "<PRICE>"
 
+# Global Variables Classification
+clf_model = None
+clf_encoder = None
+
 class TextRequest(BaseModel):
     text: str
 
-# =========================================================
-# TRIK BYPASS: BANGUN ARSITEKTUR MANUAL BERDASARKAN LOG
-# =========================================================
 def build_manual_model(vocab_size=7929, max_len=64, num_classes=8):
-    # Nama setiap layer disamakan persis dengan log agar weights menempel sempurna
     inputs = tf.keras.Input(shape=(max_len,), dtype='int32', name='tokens')
-    
     x = tf.keras.layers.Embedding(input_dim=vocab_size, output_dim=64, mask_zero=True, name='embedding')(inputs)
     x = tf.keras.layers.SpatialDropout1D(0.3, name='spatial_dropout1d')(x)
     
-    # BiLSTM 1
     lstm1_f = tf.keras.layers.LSTM(128, return_sequences=True, dropout=0.4, name='forward_lstm')
     lstm1_b = tf.keras.layers.LSTM(128, return_sequences=True, go_backwards=True, dropout=0.4, name='backward_lstm')
     x = tf.keras.layers.Bidirectional(lstm1_f, backward_layer=lstm1_b, name='bilstm_1')(x)
     
-    # BiLSTM 2
     lstm2_f = tf.keras.layers.LSTM(64, return_sequences=True, dropout=0.4, name='forward_lstm_1')
     lstm2_b = tf.keras.layers.LSTM(64, return_sequences=True, go_backwards=True, dropout=0.4, name='backward_lstm_1')
     x = tf.keras.layers.Bidirectional(lstm2_f, backward_layer=lstm2_b, name='bilstm_2')(x)
     
-    # TimeDistributed Dense 1
     dense1 = tf.keras.layers.Dense(64, activation='relu', name='dense')
     x = tf.keras.layers.TimeDistributed(dense1, name='proj')(x)
     x = tf.keras.layers.Dropout(0.4, name='dropout')(x)
     
-    # TimeDistributed Dense 2 (Logits)
     dense2 = tf.keras.layers.Dense(num_classes, activation='linear', name='dense_1')
     x = tf.keras.layers.TimeDistributed(dense2, name='logits')(x)
-    
-    # Output Softmax
     probs = tf.keras.layers.Activation('softmax', name='probabilities')(x)
     
     return tf.keras.Model(inputs=inputs, outputs=probs, name='BiLSTM_NER')
@@ -72,15 +66,18 @@ def get_ner_model():
     if ner_model is None:
         with open("vocab.pkl", "rb") as f:
             vocab = pickle.load(f)
-        
-        # 1. Buat kerangkanya
         ner_model = build_manual_model()
-        
-        # 2. HANYA isi bobotnya (Ini yang menyelamatkan kita dari error JSON!)
         ner_model.load_weights("receipt_ner_model.keras")
-        
     return ner_model, vocab
 
+def get_clf_model():
+    global clf_model, clf_encoder
+    if clf_model is None:
+        with open('penerjemah_kategori.pkl', 'rb') as file:
+            clf_encoder = pickle.load(file)
+        clf_model = tf.keras.models.load_model('model_strukly.keras', compile=False)
+    return clf_model, clf_encoder
+
 def normalize_token(token: str):
     token = token.strip()
     if re.fullmatch(r"\d{1,3}([.,]\d{3})*", token): return PRICE_TOKEN
@@ -93,47 +90,63 @@ def predict_ner(request: TextRequest):
     try:
         raw_text = request.text
         entities = {"MERCHANT": [], "ITEM": [], "TOTAL": []}
+        kategori_struk = "Tidak Terklasifikasi"
         
         if not raw_text.strip():
-            return {"success": True, "entities": entities}
+            return {"success": True, "entities": entities, "classes": kategori_struk}
 
+        # --- PROSES NER ---
         model, vcb = get_ner_model()
-        
         tokens = [t for line in raw_text.split("\n") for t in line.strip().split() if t]
-        if not tokens:
-            return {"success": True, "entities": entities}
-
-        norm_tokens = [normalize_token(t) for t in tokens]
         
-        # Hanya butuh array X, mask dihapus karena sudah otomatis ditangani layer Embedding
-        X = np.zeros((1, MAX_SEQ_LEN), dtype=np.int32)
-        length = min(len(norm_tokens), MAX_SEQ_LEN)
-        
-        for i in range(length):
-            X[0, i] = vcb.get(norm_tokens[i], vcb.get(UNK_TOKEN, 0))
+        if tokens:
+            norm_tokens = [normalize_token(t) for t in tokens]
+            X = np.zeros((1, MAX_SEQ_LEN), dtype=np.int32)
+            length = min(len(norm_tokens), MAX_SEQ_LEN)
+            
+            for i in range(length):
+                X[0, i] = vcb.get(norm_tokens[i], vcb.get(UNK_TOKEN, 0))
+                
+            probs = model(X, training=False).numpy()
+            labels = [IDX2LABEL[int(np.argmax(probs[0, i]))] for i in range(length)]
             
-        # Inference super cepat
-        probs = model(X, training=False).numpy()
-        labels = [IDX2LABEL[int(np.argmax(probs[0, i]))] for i in range(length)]
+            current_type = None
+            span = []
+            for tok, lbl in zip(tokens[:length], labels):
+                if lbl.startswith("B-"):
+                    if current_type: entities[current_type].append(" ".join(span))
+                    current_type = lbl[2:]
+                    span = [tok]
+                elif lbl.startswith("I-") and current_type == lbl[2:]:
+                    span.append(tok)
+                else:
+                    if current_type: entities[current_type].append(" ".join(span))
+                    current_type = None
+                    span = []
+            if current_type: entities[current_type].append(" ".join(span))
+
+        # --- PROSES KLASIFIKASI ---
+        items_list = entities.get("ITEM", [])
+        teks_untuk_klasifikasi = " ".join(items_list)
         
-        current_type = None
-        span = []
-        for tok, lbl in zip(tokens[:length], labels):
-            if lbl.startswith("B-"):
-                if current_type: entities[current_type].append(" ".join(span))
-                current_type = lbl[2:]
-                span = [tok]
-            elif lbl.startswith("I-") and current_type == lbl[2:]:
-                span.append(tok)
-            else:
-                if current_type: entities[current_type].append(" ".join(span))
-                current_type = None
-                span = []
-        if current_type: entities[current_type].append(" ".join(span))
+        if teks_untuk_klasifikasi:
+            try:
+                c_model, c_encoder = get_clf_model()
+                input_data = np.array([teks_untuk_klasifikasi])
+                
+                hasil_prediksi = c_model(input_data, training=False).numpy()
+                indeks_tertinggi = np.argmax(hasil_prediksi[0])
+                kategori_struk = c_encoder.inverse_transform([indeks_tertinggi])[0]
+            except Exception as e:
+                print(f"Error Klasifikasi: {e}")
 
-        return {"success": True, "entities": entities}
+        return {
+            "success": True, 
+            "entities": entities, 
+            "classes": kategori_struk 
+        }
         
     except Exception as e:
         import traceback
-        print(f"Error NER detail: {traceback.format_exc()}")
-        return {"success": False, "error": str(e), "entities": {"MERCHANT": [], "ITEM": [], "TOTAL": []}}
+        print(f"Error NER/CLF detail: {traceback.format_exc()}")
+        return {"success": False, "error": str(e), "entities": {"MERCHANT": [], "ITEM": [], "TOTAL": []}, "classes": "Error"}