Update src

app.py

@@ -1,4 +1,4 @@
-from aiofiles import os
+import os
 import gradio as gr
 import torch
 from transformers import BartForConditionalGeneration, BartTokenizer
@@ -6,17 +6,17 @@ import re
 import numpy as np
 import networkx as nx
 from typing import List, Dict
-from src.utils.get_model import get_summarizer
+from src.utils.get_model import get_summarizer, get_extractive_model, get_extractive_abstractive
 from src.preprocessing.edu_sentences import preprocess_external_text
-from src.utils.get_model import get_extractive_model
 from src.model.baseline_extractive_model import get_trigrams
-
-hf_token = os.environ.get("HF_TOKEN")
+from dotenv import load_dotenv
+load_dotenv()  # Tải biến môi trường từ file .env
+os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")  # Thay bằng token của bạn nếu cần
 
 REPO_ID_baseline_model = "Reality8081/bart-base"
 REPO_ID_baseline_model_edu = "Reality8081/bart-base-edu"
-REPO_ID_baseline_extractive_model = "Reality8081/bart-extractive"
-REPO_ID_baseline_extractive_model_edu = "Reality8081/bart-extractive-edu"
+REPO_ID_baseline_extractive_model = "Reality8081/bart_extractive"
+REPO_ID_baseline_extractive_model_edu = "Reality8081/bart_extractive-edu"
 REPO_ID_Extabs_model = "Reality8081/bart-encoder-decoder"
 REPO_ID_Extabs_model_edu = "Reality8081/bart-encoder-decoder-edu"
 
@@ -52,12 +52,14 @@ def model_baseline_extractive(prepro_dict: Dict, top_n = 5) -> str:
         repo_id = REPO_ID_baseline_extractive_model_edu
     else:
         repo_id = REPO_ID_baseline_extractive_model        
-    model = get_summarizer(repo_id=repo_id, base_model_name="facebook/bart-large", device=device)
-    with torch.no_grad():
+    model = get_extractive_model(repo_id=repo_id, device=device)
+    model = model.to(torch.float32)
+    model.eval()
+    with torch.no_grad(), torch.autocast(device_type="cuda" if torch.cuda.is_available() else "cpu"):
         outputs = model(input_ids=input_ids, attention_mask=attention_mask)
         # Sử dụng Sigmoid đưa logit về khoảng (0, 1) để lấy xác suất
-        probs = torch.sigmoid(outputs['logits']).squeeze(0).cpu().numpy()
-        
+        logits = outputs['logits'].to(torch.float32)
+        probs = torch.sigmoid(logits).squeeze(0).cpu().numpy()        
     segment_scores = []
     current_idx = 1 # Bỏ qua token đặc biệt <s> ở đầu chuỗi
     
@@ -109,7 +111,7 @@ def model_extractive_abstract(prepro_dict: Dict) -> str:
         repo_id = REPO_ID_Extabs_model_edu
     else:
         repo_id = REPO_ID_Extabs_model 
-    model = get_summarizer(repo_id=repo_id, base_model_name="facebook/bart-large", device=device)
+    model = get_extractive_abstractive(repo_id=repo_id, base_model_name="facebook/bart-large", device=device)
     with torch.no_grad():
         summary_ids = model.generate_summary(
             input_ids=input_ids,
@@ -140,7 +142,7 @@ def ATS(
         prepro_dict = preprocess_external_text(text, reference_summary, segmentation_method='edu')
 
     # Step 2: Chọn model
-    if model == "Baseline Model: TextRank + Vanilla BART":
+    if model == "Baseline Model":
         result = model_baseline(prepro_dict)
     elif model == "Baseline Model with Extractive":
         result = model_baseline_extractive(prepro_dict)
@@ -153,11 +155,7 @@ def ATS(
 # ====================== GRADIO INTERFACE ======================
 with gr.Blocks(
     title="Automated Text Summarization System",
-    theme=gr.themes.Soft(),
-    css="""
-    .gradio-container {max-width: 1200px; margin: auto;}
-    .title {text-align: center; margin-bottom: 10px;}
-    """
+    
 ) as demo:
     gr.Markdown(
         """
@@ -173,7 +171,6 @@ with gr.Blocks(
                 placeholder="Paste your long text here (up to several thousand words)...",
                 lines=12,
                 max_lines=30,
-                show_copy_button=True
             )
 
         with gr.Column(scale=1):
@@ -211,7 +208,6 @@ with gr.Blocks(
         label="📄 Summary Result",
         lines=10,
         placeholder="The result will appear here...",
-        show_copy_button=True
     )
 
     # Connect button click
@@ -245,4 +241,8 @@ with gr.Blocks(
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(theme=gr.themes.Soft(),
+    css="""
+    .gradio-container {max-width: 1200px; margin: auto;}
+    .title {text-align: center; margin-bottom: 10px;}
+    """)

src/model/baseline_extractive_model.py

@@ -2,40 +2,53 @@ import torch
 import torch.nn as nn
 import numpy as np
 from transformers import BartModel, BartTokenizer
+from huggingface_hub import PyTorchModelHubMixin
 
-class BartExtractiveSummarizer(nn.Module):
+class BartExtractiveSummarizer(nn.Module, PyTorchModelHubMixin):
     def __init__(self, model_name="facebook/bart-large"):
         super(BartExtractiveSummarizer, self).__init__()
         self.encoder = BartModel.from_pretrained(model_name).encoder
         hidden_size = self.encoder.config.hidden_size
         self.classifier = nn.Linear(hidden_size, 1)
 
+        # Force float32 from the beginning
+        self.to(torch.float32)
+
     def forward(self, input_ids, attention_mask, saliency_mask=None, **kwargs):
+        device = next(self.parameters()).device
+
+        input_ids = input_ids.to(torch.long).to(device)
+        attention_mask = attention_mask.to(torch.long).to(device)
+        if saliency_mask is not None:
+            saliency_mask = saliency_mask.to(torch.float32).to(device)
+
+        # Extra safety: ensure encoder stays in float32
+        if self.encoder.parameters().__next__().dtype != torch.float32:
+            self.encoder = self.encoder.to(torch.float32)
+
         encoder_outputs = self.encoder(
             input_ids=input_ids,
             attention_mask=attention_mask
         )
-        hidden_states = encoder_outputs.last_hidden_state 
-        logits = self.classifier(hidden_states).squeeze(-1) 
         
+        hidden_states = encoder_outputs.last_hidden_state.float()
+        logits = self.classifier(hidden_states).squeeze(-1)
+
         loss = None
         if saliency_mask is not None:
             active_loss = attention_mask.view(-1) == 1
             active_logits = logits.view(-1)[active_loss]
             active_labels = saliency_mask.view(-1)[active_loss].float()
-            
-            # --- TỐI ƯU: TỰ ĐỘNG TÍNH CLASS WEIGHT CHO TỪNG BATCH ---
+
             num_pos = active_labels.sum()
             num_neg = active_labels.size(0) - num_pos
-            
-            if num_pos > 0:
-                weight = torch.clamp(num_neg / num_pos, max=10.0) 
-            else:
-                weight = torch.tensor(1.0).to(logits.device)
-                
+
+            weight = torch.tensor(num_neg / num_pos if num_pos > 0 else 1.0, 
+                                dtype=torch.float32, device=logits.device)
+
             loss_fct = nn.BCEWithLogitsLoss(pos_weight=weight)
             loss = loss_fct(active_logits, active_labels)
-            
+
         return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}
     
 def get_trigrams(text: str):

src/model/extabs.py

@@ -1,8 +1,9 @@
 import torch
 import torch.nn as nn
 from transformers import BartForConditionalGeneration, BartTokenizer
+from huggingface_hub import PyTorchModelHubMixin
 
-class EXTABSModel(nn.Module):
+class EXTABSModel(nn.Module, PyTorchModelHubMixin):
     def __init__(self, model_name="facebook/bart-large"):
         super(EXTABSModel, self).__init__()
         # Load kiến trúc BART gốc

src/utils/get_model.py

@@ -4,7 +4,7 @@ import torch
 from huggingface_hub import hf_hub_download
 from src.model.extabs import EXTABSModel
 import gc
-
+import os
 from safetensors.torch import load_file
 
 active_model_info = {
@@ -42,14 +42,10 @@ def get_extractive_model(repo_id: str, base_model_name: str = "facebook/bart-lar
         print(f"Đang tải mô hình Extractive từ repo: {repo_id}...")
         
         # Khởi tạo khung kiến trúc trống
-        model = BartExtractiveSummarizer(model_name=base_model_name)
-        
+        model = BartExtractiveSummarizer.from_pretrained(repo_id, model_name=base_model_name)        
         # Sử dụng hf_hub_download để kéo file trọng số về local cache
-        model_path = hf_hub_download(repo_id=repo_id, filename="model_state.bin")
-        state_dict = load_file(model_path)
-        # Load trọng số vào model
-        model.load_state_dict(state_dict)
-        
+        # Load trọng số vào model  
+        model = model.to(torch.float32)      
         model.to(device)
         model.eval() # Chuyển mô hình sang chế độ inference
         
@@ -57,21 +53,18 @@ def get_extractive_model(repo_id: str, base_model_name: str = "facebook/bart-lar
         active_model_info["repo_id"] = repo_id        
     return active_model_info["model"]
 
-def get_extractive_abstractive(repo_id: str, base_model_name: str = "facebook/bart-large", device: torch.device = "cpu"):
+def get_extractive_abstractive(repo_id: str,base_model_name: str = "facebook/bart-large", device: torch.device = "cpu"):
     """Tải và lưu cache mô hình Custom Extractive từ Hugging Face Hub"""
     if active_model_info["repo_id"] != repo_id:
         clear_memory() # XÓA SẠCH MODEL CŨ
         print(f"Đang tải mô hình Extractive từ repo: {repo_id}...")
         
         # Khởi tạo khung kiến trúc trống
-        model = EXTABSModel(model_name=base_model_name)
+        model = EXTABSModel.from_pretrained(repo_id, model_name=base_model_name)   
         
         # Sử dụng hf_hub_download để kéo file trọng số về local cache
-        model_path = hf_hub_download(repo_id=repo_id, filename="model_state.bin")
-        state_dict = load_file(model_path)
         # Load trọng số vào model
-        model.load_state_dict(state_dict)
-        
+        model = model.to(torch.float32)        
         model.to(device)
         model.eval() # Chuyển mô hình sang chế độ inference