Spaces:
Sleeping
Sleeping
| import networkx as nx | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| import torch | |
| from transformers import BartForConditionalGeneration, BartTokenizer | |
| def textrank_summarize(sentences, top_n=3): | |
| """ | |
| Trích xuất câu quan trọng bằng TextRank + TF-IDF. | |
| Đầu vào 'sentences' là một list các câu (hoặc EDUs) trong một văn bản. | |
| """ | |
| # Xử lý trường hợp bài báo quá ngắn | |
| if len(sentences) <= top_n: | |
| return " ".join(sentences) | |
| try: | |
| # Bước 1: Khởi tạo TfidfVectorizer và fit_transform tập sentences của 1 bài báo | |
| vectorizer = TfidfVectorizer(stop_words='english') | |
| tfidf_matrix = vectorizer.fit_transform(sentences) | |
| # Bước 2: Tính ma trận tương đồng Cosine | |
| similarity_matrix = cosine_similarity(tfidf_matrix, tfidf_matrix) | |
| # Bước 3: Đưa ma trận vào networkx tạo đồ thị và tính PageRank | |
| nx_graph = nx.from_numpy_array(similarity_matrix) | |
| scores = nx.pagerank(nx_graph) | |
| # Bước 4: Sắp xếp điểm số và chọn top_n câu | |
| ranked_sentences = sorted(((scores[i], s) for i, s in enumerate(sentences)), reverse=True) | |
| # Giữ đúng thứ tự xuất hiện của câu trong văn bản gốc để dễ đọc | |
| top_sentences_indices = sorted([sentences.index(ranked_sentences[i][1]) for i in range(top_n)]) | |
| summary = " ".join([sentences[i] for i in top_sentences_indices]) | |
| return summary | |
| except Exception as e: | |
| # Fallback về Lead-N nếu đồ thị lỗi (do câu rỗng hoặc không có từ vựng) | |
| return " ".join(sentences[:top_n]) | |
| class BartSummarizer: | |
| def __init__(self, model_path="facebook/bart-base"): | |
| """ | |
| Khởi tạo mô hình và tokenizer. | |
| model_path có thể là repo trên Hugging Face hoặc đường dẫn local chứa weights. | |
| """ | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Loading BART model from '{model_path}' onto {self.device}...") | |
| self.tokenizer = BartTokenizer.from_pretrained(model_path) | |
| self.model = BartForConditionalGeneration.from_pretrained(model_path) | |
| self.model.to(self.device) | |
| self.model.eval() # Chuyển sang chế độ inference ngay từ đầu | |
| def summarize(self, text, max_input_length=512, max_output_length=128, min_output_length=30): | |
| """ | |
| Hàm sinh tóm tắt cho một đoạn văn bản đầu vào. | |
| """ | |
| with torch.no_grad(): | |
| # Cắt ngắn đầu vào để chống quá tải GPU, đồng bộ với lúc train | |
| inputs = self.tokenizer( | |
| text, | |
| max_length=max_input_length, | |
| truncation=True, | |
| padding=True, | |
| return_tensors="pt" | |
| ).to(self.device) | |
| # Sinh văn bản tóm tắt | |
| summary_ids = self.model.generate( | |
| input_ids=inputs["input_ids"], | |
| attention_mask=inputs["attention_mask"], | |
| max_length=max_output_length, | |
| min_length=min_output_length, | |
| num_beams=4, | |
| length_penalty=2.0, # Ưu tiên sinh câu trọn vẹn | |
| no_repeat_ngram_size=3, # Chống ảo giác, lặp từ | |
| early_stopping=True | |
| ) | |
| # Decode kết quả về dạng text | |
| summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True) | |
| return summary | |
| # Cách gọi trong app: | |
| # summarizer = BartSummarizer("duong_dan_model_cua_ban_tren_huggingface") | |
| # result = summarizer.summarize("Đoạn văn bản cần tóm tắt...") |