Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import streamlit as st | |
| from transformers import BartTokenizer, BartForConditionalGeneration | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.cluster import KMeans | |
| from nltk.tokenize import sent_tokenize | |
| import nltk | |
| # Download required NLTK data | |
| def download_nltk_data(): | |
| try: | |
| nltk.data.find('tokenizers/punkt') | |
| except LookupError: | |
| nltk.download('punkt', quiet=True) | |
| nltk.download('punkt_tab', quiet=True) | |
| download_nltk_data() | |
| # Load BART model and tokenizer | |
| def load_model(): | |
| tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn') | |
| model = BartForConditionalGeneration.from_pretrained('facebook/bart-large-cnn') | |
| return tokenizer, model | |
| tokenizer, model = load_model() | |
| # Function for single document summarization | |
| def summarize(text, max_length=150, min_length=50): | |
| inputs = tokenizer([text], max_length=1024, truncation=True, return_tensors='pt') | |
| summary_ids = model.generate( | |
| inputs['input_ids'], | |
| num_beams=4, | |
| max_length=max_length, | |
| min_length=min_length, | |
| early_stopping=True | |
| ) | |
| return tokenizer.decode(summary_ids[0], skip_special_tokens=True) | |
| # Function for clustering documents | |
| def cluster_documents(documents, n_clusters=3): | |
| vectorizer = TfidfVectorizer(stop_words='english') | |
| X = vectorizer.fit_transform(documents) | |
| kmeans = KMeans(n_clusters=n_clusters, random_state=42).fit(X) | |
| return kmeans.labels_ | |
| # Function to split long text into chunks | |
| def chunk_text(text, max_words=1000): | |
| if len(text.split()) <= max_words: | |
| return [text] | |
| sentences = sent_tokenize(text) | |
| current_chunk = [] | |
| chunks = [] | |
| current_word_count = 0 | |
| for sentence in sentences: | |
| sentence_words = len(sentence.split()) | |
| if current_word_count + sentence_words <= max_words: | |
| current_chunk.append(sentence) | |
| current_word_count += sentence_words | |
| else: | |
| chunks.append(" ".join(current_chunk)) | |
| current_chunk = [sentence] | |
| current_word_count = sentence_words | |
| if current_chunk: | |
| chunks.append(" ".join(current_chunk)) | |
| return chunks | |
| # Function for multi-document summarization | |
| def multi_document_summarize(documents): | |
| results = { | |
| 'individual_summaries': [], | |
| 'cluster_summaries': [], | |
| 'final_summary': None | |
| } | |
| # Generate individual document summaries | |
| for doc in documents: | |
| chunks = chunk_text(doc) | |
| doc_summary = " ".join([summarize(chunk) for chunk in chunks]) | |
| results['individual_summaries'].append(doc_summary) | |
| # Cluster documents if we have multiple | |
| if len(documents) > 1: | |
| n_clusters = min(3, len(documents)) | |
| clusters = cluster_documents(documents, n_clusters=n_clusters) | |
| # Generate cluster summaries | |
| for cluster_id in np.unique(clusters): | |
| cluster_docs = [doc for doc, c in zip(documents, clusters) if c == cluster_id] | |
| combined_text = " ".join(cluster_docs) | |
| chunks = chunk_text(combined_text) | |
| cluster_summary = " ".join([summarize(chunk) for chunk in chunks]) | |
| results['cluster_summaries'].append({ | |
| 'doc_indices': [i for i, c in enumerate(clusters) if c == cluster_id], | |
| 'summary': cluster_summary | |
| }) | |
| # Generate final comprehensive summary | |
| all_summaries = results['individual_summaries'] + [cs['summary'] for cs in results['cluster_summaries']] | |
| results['final_summary'] = summarize(" ".join(all_summaries), max_length=200, min_length=100) | |
| else: | |
| results['final_summary'] = results['individual_summaries'][0] | |
| return results | |
| # Streamlit UI | |
| st.title("Multi-Document Summarization") | |
| st.write("This system summarizes multiple related documents using clustering and BART.") | |
| uploaded_files = st.file_uploader( | |
| "Upload text files (one document per file)", | |
| type=['txt'], | |
| accept_multiple_files=True | |
| ) | |
| if uploaded_files: | |
| documents = [] | |
| for file in uploaded_files: | |
| documents.append(file.read().decode("utf-8")) | |
| if st.button("Generate Summary"): | |
| with st.spinner("Processing documents..."): | |
| try: | |
| results = multi_document_summarize(documents) | |
| # Display individual document summaries | |
| st.subheader("Individual Document Summaries") | |
| for i, (doc, summary) in enumerate(zip(documents, results['individual_summaries'])): | |
| with st.expander(f"Document {i+1} Summary (First 200 chars: {doc[:200]}...)"): | |
| st.write(summary) | |
| # Display cluster summaries if we have multiple documents | |
| if len(documents) > 1: | |
| st.subheader("Cluster Summaries") | |
| for i, cluster in enumerate(results['cluster_summaries']): | |
| with st.expander(f"Cluster {i+1} Summary (Documents: {', '.join(str(x+1) for x in cluster['doc_indices'])}))"): | |
| st.write(cluster['summary']) | |
| # Display final summary | |
| st.subheader("Final Comprehensive Summary") | |
| st.write(results['final_summary']) | |
| except Exception as e: | |
| st.error(f"An error occurred: {str(e)}") | |
| st.info("Please try again or upload different documents.") | |
| else: | |
| st.info("Please upload some text files to get started.") |