add infer vit5

app.py

@@ -1,8 +1,14 @@
 import streamlit as st
 from io import StringIO
 from chdg_inference import infer
+from infer_concat import vit5_infer
 
-st.title("Tóm tắt đa văn bản tiếng Việt")
+st.set_page_config(layout="wide")
+st.title("Tóm tắt Đa văn bản Tiếng Việt")
+
+col1, col2 = st.columns([1, 1])
+col2_title, = col2.columns(1)
+col2_chdg, col2_vit5 = col2.columns(2)
 
 # Initialize session state
 if 'num_docs' not in st.session_state:
@@ -14,40 +20,32 @@ if 'docs' not in st.session_state:
 def add_text_area():
     st.session_state.num_docs += 1
 
-
 # Button to add a new text area
-st.button("Thêm văn bản", on_click=add_text_area)
+col1.button("Thêm văn bản", on_click=add_text_area)
 
 # Display text areas for document input
 for i in range(st.session_state.num_docs):
-    doc = st.text_area(f"Văn bản {i+1}", key=f"doc_{i}", height=200)
+    doc = col1.text_area(f"Văn bản {i+1}", key=f"doc_{i}", height=150)
     doc.replace('\r', '\n')
+    doc.replace('\"', "'")
     if len(st.session_state.docs) <= i:
         st.session_state.docs.append(doc)
     else:
         st.session_state.docs[i] = doc
-# Display the documents for verification
-# st.write("**Entered Documents:**")
-# st.write(st.session_state.docs)
-# uploaded_file = st.file_uploader(label="Chọn file văn bản")
 
-category = st.selectbox("Chọn chủ để của văn bản: ", ['Giáo dục', 'Giải trí - Thể thao', 'Khoa học - Công nghệ', 'Kinh tế', 'Pháp luật', 'Thế giới', 'Văn hóa - Xã hội', 'Đời sống'])
+category = col1.selectbox("Chọn chủ để của văn bản: ", ['Giáo dục', 'Giải trí - Thể thao', 'Khoa học - Công nghệ', 'Kinh tế', 'Pháp luật', 'Thế giới', 'Văn hóa - Xã hội', 'Đời sống'])
 
 def summarize():
-    # if uploaded_file is not None:
-    #     stringio = StringIO(uploaded_file.getvalue().decode("utf-8"))
-    #     full_text = stringio.read()
-    #     summ, docs = infer(full_text, category)
-    #     st.subheader("Kết quả: ")
-    #     st.write(summ)
-    #     st.subheader("Docs: ")
-    #     st.write(docs)
-    # else:
-    #     st.error("Hãy tải file văn bản lên")
-    summ, docs = infer(st.session_state.docs, category)
-    st.subheader("Kết quả")
-    st.write(summ)
-    st.write(docs)
-
-if st.button("Tóm tắt"):
+    summ, _ = infer(st.session_state.docs, category)
+    with col2.container():
+        col2_title.subheader("Kết quả: ")
+        col2_title.write("\n")
+
+    with col2.container():
+        col2_chdg.write("CHDG:")
+        col2_chdg.write(summ)
+        summ_vit5 = vit5_infer(st.session_state.docs)
+        col2_vit5.write(summ_vit5)
+    
+if col1.button("Tóm tắt"):
     summarize()

infer_concat.py

@@ -0,0 +1,109 @@
+# create dataset class
+from torch.utils.data import Dataset, DataLoader
+import torch
+import json
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import time
+
+
+class Dataset4Summarization(Dataset):
+	def __init__(self, data, tokenizer, max_length=1024*3, chunk_length =1024):
+		self.data = data
+		self.tokenizer = tokenizer
+		self.max_length = max_length
+		self.chunk_length = chunk_length
+
+	def __len__(self):
+		return len(self.data)
+	
+	def chunking(self, text):
+		chunks = []
+		for i in range(0, self.max_length, self.chunk_length):
+			chunks.append(text[i:i+self.chunk_length])
+		return chunks
+
+	def __getitem__(self, idx):
+		sample = self.data[idx]
+		inputs = self.tokenizer(sample, return_tensors='pt', padding='max_length', truncation=True, max_length=self.max_length)
+
+		list_chunk = self.chunking(inputs['input_ids'].squeeze())
+		list_attention_mask = self.chunking(inputs['attention_mask'].squeeze())
+
+
+		return {
+			'list_input_ids': list_chunk,
+			'list_att_mask' : list_attention_mask,
+		}
+	
+
+def process_data_infer(data):
+	single_documents = data.get('single_documents', [])
+
+	
+	result = []
+	for doc in single_documents:
+		raw_text = doc.get('raw_text', '')
+		result.append(raw_text)
+
+	return " ".join(result)
+
+
+def processing_data_infer(input_file):
+	all_results = []
+	
+	with open(input_file, 'r', encoding='utf-8') as file:
+		for line in file:
+			data = json.loads(line.strip())
+			result = process_data_infer(data)
+			all_results.append(result)
+
+	return all_results
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base-vietnews-summarization")
+model = AutoModelForSeq2SeqLM.from_pretrained("VietAI/vit5-base-vietnews-summarization")
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model.to(device)
+
+model.load_state_dict(torch.load("./weight_cp19_model.pth", map_location=torch.device('cpu')))
+
+# For other demo purpose, you just need to make sure data is list of documents [document1, document2]
+
+# batch_size need to be 1,
+@torch.no_grad()
+def infer_2_hier(model, data_loader, device, tokenizer):
+    model.eval()
+    start = time.time()
+    all_summaries = []
+    for iter in data_loader:
+        summaries = []
+        inputs = iter['list_input_ids']
+        att_mask = iter['list_att_mask']
+        
+        for i in range(len(inputs)):
+            # Check if the input tensor is all zeros
+            if torch.all(inputs[i] == 0):
+                # If the input is all zeros, skip this iteration
+                continue
+            else:
+                summary = model.generate(inputs[i].to(device),
+                                         attention_mask=att_mask[i].to(device),
+                                         max_length=128,
+                                         num_beams=12,
+                                         num_return_sequences=1)
+                summaries.append(summary)
+        summaries = torch.cat(summaries, dim = 1)
+        for k in summaries:
+                all_summaries.append(tokenizer.decode(k, skip_special_tokens=True))
+
+    
+    end = time.time()
+    print(f"Time: {end-start}")
+    return all_summaries
+
+def vit5_infer(data):
+	dataset = Dataset4Summarization(data, tokenizer)
+	data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=2)
+	result = infer_2_hier(model, data_loader, device, tokenizer)
+	return result

requirements.txt

@@ -1,7 +1,7 @@
-torch
-rouge
-transformers
-underthesea
-numpy
-pandas
-scikit-learn
+torch==2.1.2
+rouge==1.0.1
+transformers==4.39.2
+underthesea==6.8.4
+numpy==1.25.1
+pandas==2.1.1
+scikit-learn==1.3.0