summarize / summarizer.py
Ayush
initial commit - AI summarizer app
0001c2f
Raw
History Blame Contribute Delete
1.81 kB
from transformers import pipeline
import re
from functools import lru_cache
MODELS = {
"BART (facebook/bart-large-cnn)": "facebook/bart-large-cnn",
"T5 Large (t5-large)": "t5-large",
"Pegasus (google/pegasus-xsum)": "google/pegasus-xsum",
}
@lru_cache(maxsize=3)
def load_model(name):
print(f"Loading model → {name}")
return pipeline("summarization", model=name, framework="pt")
def clean_text(text):
return re.sub(r'\s+', ' ', text).strip()
def summarize_once(text, model_name, max_length, min_length):
model = load_model(model_name)
text = clean_text(text)[:2000]
summary = model(text, max_length=max_length, min_length=min_length, do_sample=False)
return summary[0]["summary_text"]
def chunk_text(text, size=300):
words = text.split()
for i in range(0, len(words), size):
yield " ".join(words[i:i+size])
# Normal Summarization (NO final compression)
def summarize_basic(text, model_name, max_length, min_length, chunk_size=300, progress=None):
text = clean_text(text)
chunks = list(chunk_text(text, chunk_size))
if len(chunks) == 1:
if progress: progress(100)
return summarize_once(text, model_name, max_length, min_length)
summaries=[]
for i,chunk in enumerate(chunks):
summaries.append(summarize_once(chunk, model_name, max_length, min_length))
if progress: progress(int(((i+1)/len(chunks))*100))
return "\n\n".join(summaries) # <-- no merge summarization
# Extra Summarization (Hierarchical)
def summarize_advanced(text, model_name, max_length, min_length, chunk_size=300, progress=None):
basic = summarize_basic(text, model_name, max_length, min_length, chunk_size, progress)
return summarize_once(basic, model_name, max_length//2, min_length//2) # more compression