Update py/abstractive.py

py/abstractive.py

@@ -5,8 +5,8 @@ import re
 import pickle
 import unicodedata
 import numpy as np
-import torch
-from transformers import AutoTokenizer, AutoModel
+import tensorflow as tf # Dùng TF hoàn toàn
+from transformers import AutoTokenizer, TFAutoModel # Dùng TFAutoModel thay vì AutoModel
 
 # -------------------------------------------------------------------
 # CONFIG
@@ -21,7 +21,7 @@ CONFIG = {
     "FF_DIM": 2048,
     "DROPOUT": 0.2,
     "TOKENIZER_FILE": "decoder_tokenizer_re.pkl",
-    "WEIGHTS_FILE": "decoder_only.weights.h5" 
+    "WEIGHTS_FILE": "best_model.weights.h5" # Đã đổi tên khớp với file train của bạn
 }
 
 # -------------------------------------------------------------------
@@ -38,81 +38,61 @@ def clean_text_inference(text: str) -> str:
     return text
 
 # -------------------------------------------------------------------
-# PHOBERT ENCODER (PyTorch) - ĐÃ SỬA LỖI META TENSOR TẠI ĐÂY
+# PHOBERT ENCODER (TENSORFLOW VERSION)
 # -------------------------------------------------------------------
-class PhoBERTEncoderTorch:
+class PhoBERTEncoderTF:
     def __init__(self):
-        print(">>> Loading PhoBERT (PyTorch)...")
-        self.device = torch.device("cpu")
-        
-        # === FIX: Thêm low_cpu_mem_usage=False ===
-        # Lỗi "meta tensor" xảy ra do thư viện accelerate cố gắng tối ưu RAM bằng cách tạo model rỗng.
-        # low_cpu_mem_usage=False buộc model load đầy đủ weights vào RAM ngay lập tức.
-        try:
-            self.model = AutoModel.from_pretrained("vinai/phobert-base", low_cpu_mem_usage=False).to(self.device)
-        except TypeError:
-            # Fallback nếu phiên bản transformers cũ không hỗ trợ tham số này
-            self.model = AutoModel.from_pretrained("vinai/phobert-base").to(self.device)
-            
-        self.model.eval()
+        print(">>> Loading PhoBERT (TensorFlow)...")
+        # Load đúng model TF giống lúc train
+        self.model = TFAutoModel.from_pretrained("vinai/phobert-base")
         print(">>> PhoBERT loaded successfully.")
 
     def encode(self, input_ids, attention_mask):
-        with torch.no_grad():
-            ids = torch.tensor(input_ids).to(self.device)
-            mask = torch.tensor(attention_mask).to(self.device)
-            outputs = self.model(ids, attention_mask=mask)
-            return outputs.last_hidden_state.cpu().numpy()
-
-# -------------------------------------------------------------------
-# SAFE IMPORT TENSORFLOW
-# -------------------------------------------------------------------
-TF_AVAILABLE = True
-try:
-    import tensorflow as tf
-except Exception as e:
-    TF_AVAILABLE = False
-    _TF_ERR = e
-    tf = None
+        # Input là numpy array, TFAutoModel nhận trực tiếp
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        # Trả về last_hidden_state dạng numpy hoặc tensor
+        return outputs.last_hidden_state
 
 # -------------------------------------------------------------------
 # DECODER TRANSFORMER (TensorFlow)
 # -------------------------------------------------------------------
-if TF_AVAILABLE:
-    class TransformerDecoderBlock(tf.keras.layers.Layer):
-        def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):
-            super().__init__(**kwargs)
-            self.att1 = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
-            self.att2 = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
-            self.ffn = tf.keras.Sequential([
-                tf.keras.layers.Dense(ff_dim, activation="relu"),
-                tf.keras.layers.Dense(embed_dim),
-            ])
-            self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
-            self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
-            self.ln3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
-            self.drop1 = tf.keras.layers.Dropout(rate)
-            self.drop2 = tf.keras.layers.Dropout(rate)
-            self.drop3 = tf.keras.layers.Dropout(rate)
-
-        def call(self, x, enc_output, training=None):
-            attn1 = self.att1(x, x, use_causal_mask=True)
-            out1 = self.ln1(x + self.drop1(attn1, training=training))
-            attn2 = self.att2(out1, enc_output)
-            out2 = self.ln2(out1 + self.drop2(attn2, training=training))
-            ffn_out = self.ffn(out2)
-            return self.ln3(out2 + self.drop3(ffn_out, training=training))
+class TransformerDecoderBlock(tf.keras.layers.Layer):
+    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.att1 = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
+        self.att2 = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
+        self.ffn = tf.keras.Sequential([
+            tf.keras.layers.Dense(ff_dim, activation="relu"),
+            tf.keras.layers.Dense(embed_dim),
+        ])
+        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+        self.ln3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+        self.drop1 = tf.keras.layers.Dropout(rate)
+        self.drop2 = tf.keras.layers.Dropout(rate)
+        self.drop3 = tf.keras.layers.Dropout(rate)
+
+    def call(self, x, enc_output, training=None):
+        attn1 = self.att1(x, x, use_causal_mask=True)
+        out1 = self.ln1(x + self.drop1(attn1, training=training))
+        attn2 = self.att2(out1, enc_output)
+        out2 = self.ln2(out1 + self.drop2(attn2, training=training))
+        ffn_out = self.ffn(out2)
+        return self.ln3(out2 + self.drop3(ffn_out, training=training))
 
 # -------------------------------------------------------------------
-# BUILD DECODER MODEL
+# BUILD INFERENCE MODEL
 # -------------------------------------------------------------------
 def build_inference_model():
+    # 1. Inputs
     enc_raw_input = tf.keras.Input(shape=(None, 768), name='enc_raw_input')
     dec_inputs_inf = tf.keras.Input(shape=(None,), dtype=tf.int32, name='dec_inputs_inf')
 
+    # 2. Projection Layer (Tên layer phải khớp file train: encoder_projection)
     enc_out = tf.keras.layers.Dense(CONFIG["EMBED_DIM"], activation="linear", name="encoder_projection")(enc_raw_input)
     enc_out = tf.keras.layers.Dropout(CONFIG["DROPOUT"], name="encoder_dropout")(enc_out)
 
+    # 3. Embeddings
     dec_token_emb = tf.keras.layers.Embedding(CONFIG["MAX_VOCAB_OUT"], CONFIG["EMBED_DIM"], mask_zero=True, name='dec_token_emb')
     dec_pos_emb = tf.keras.layers.Embedding(CONFIG["MAX_SUMMARY_LEN"], CONFIG["EMBED_DIM"], name='dec_pos_emb')
 
@@ -126,6 +106,7 @@ def build_inference_model():
 
     dec_emb_inf = tf.keras.layers.Lambda(add_pos_emb_inf, name='dec_emb_plus_pos_inf')(dec_inputs_inf)
 
+    # 4. Decoder Blocks
     dec_out = dec_emb_inf
     for i in range(CONFIG["NUM_LAYERS"]):
         block = TransformerDecoderBlock(
@@ -137,6 +118,7 @@ def build_inference_model():
         )
         dec_out = block(dec_out, enc_out, training=False)
 
+    # 5. Output
     outputs_inf = tf.keras.layers.Dense(CONFIG["MAX_VOCAB_OUT"], activation='softmax', name='output_dense')(dec_out)
 
     model = tf.keras.Model(inputs=[enc_raw_input, dec_inputs_inf], outputs=outputs_inf, name="inference_decoder_export")
@@ -147,12 +129,10 @@ def build_inference_model():
 # -------------------------------------------------------------------
 class AbstractiveSummarizer:
     def __init__(self, model_dir="./models"):
-        if not TF_AVAILABLE:
-            raise RuntimeError(f"TensorFlow không khả dụng: {_TF_ERR}")
-
         self.model_dir = model_dir
         
-        self.phobert = PhoBERTEncoderTorch()
+        # Load PhoBERT (TensorFlow)
+        self.phobert_encoder = PhoBERTEncoderTF()
         self.phobert_tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base")
         
         self._load_resources()
@@ -174,8 +154,12 @@ class AbstractiveSummarizer:
         
         print(f"📥 Loading weights from {weights_path}...")
         try:
-            self.decoder_model.load_weights(weights_path)
-            print("✅ Weights loaded successfully!")
+            # Load weights. 
+            # Quan trọng: Nếu file train lưu "best_model.weights.h5", load vào cấu trúc inference này
+            # sẽ bỏ qua các layer của PhoBERT (vì inference model bắt đầu từ enc_raw_input)
+            # và chỉ load weights của Projection + Decoder.
+            self.decoder_model.load_weights(weights_path, skip_mismatch=True)
+            print("✅ Weights loaded successfully (with skip_mismatch=True for safety)!")
         except Exception as e:
             print(f"❌ Error loading weights: {e}")
 
@@ -196,6 +180,7 @@ class AbstractiveSummarizer:
                     continue
                 
                 dec_inp = np.array([seq])
+                # Dự đoán
                 preds = self.decoder_model.predict([enc_out, dec_inp], verbose=0)
                 last_token_probs = preds[0, -1, :]
                 
@@ -216,17 +201,26 @@ class AbstractiveSummarizer:
         return summary
 
     def summarize_debug(self, text, k=3):
+        # 1. Clean
         text_clean = clean_text_inference(text)
+        
+        # 2. Tokenize Input
         inp = self.phobert_tokenizer(
             [text_clean],
             max_length=CONFIG["MAX_TEXT_LEN"],
             truncation=True, padding='max_length',
-            return_tensors='np'
+            return_tensors='np' # Quan trọng: trả về numpy cho TF
         )
-        enc_out = self.phobert.encode(inp['input_ids'], inp['attention_mask'])
-        seq = self.beam_search(enc_out, k=k)
         
-        decoded_text = self.tokenizer.sequences_to_texts([seq])[0]
+        # 3. Encode (Qua PhoBERT TF)
+        # inp['input_ids'] là numpy array
+        enc_out = self.phobert_encoder.encode(inp['input_ids'], inp['attention_mask'])
+        
+        # 4. Beam Search
+        seq_ids = self.beam_search(enc_out, k=k)
+        
+        # 5. Decode
+        decoded_text = self.tokenizer.sequences_to_texts([seq_ids])[0]
         summary = decoded_text.replace('startseq', '').replace('endseq', '').strip()
         
-        return summary, seq
+        return summary, seq_ids