CoolFace
Apppublic

bjong/blessed_Text_summarization_and_lingual_model

sourceHugging Faceafl-3.0updated 2y agoView on Hugging Face
0likes
summarizer.py78 linesDownload Raw Back to root
1from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
2from PyPDF2 import PdfReader
3from nltk.tokenize import sent_tokenize
4import torch
5import streamlit as st
6
7@st.cache_resource
8def load_model():
9    model_path = './model/'  # Update this with the correct path to your saved model
10    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
11    tokenizer = AutoTokenizer.from_pretrained(model_path)
12    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13    model.to(device)
14    return model, tokenizer, device
15
16summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
17
18def summarize_text(text, model, tokenizer, device, max_chunk_size=1024):
19    if len(text) > max_chunk_size:
20        sentences = sent_tokenize(text)
21        chunks = [sentences[i:i+10] for i in range(0, len(sentences), 10)]
22    else:
23        chunks = [text]
24
25    summaries = []
26
27    for chunk in chunks:
28        chunk_text = " ".join(chunk)
29        inputs = tokenizer(chunk_text, return_tensors="pt", max_length=1024, truncation=True)
30        input_ids = inputs["input_ids"].to(device)
31        attention_mask = inputs["attention_mask"].to(device)
32        try:
33            summary_ids = model.generate(input_ids=input_ids, attention_mask=attention_mask, length_penalty=0.8, num_beams=8, min_length=60, max_length=500)
34            summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
35            summaries.append(summary)
36        except IndexError:
37            return "Input size exceeds the maximum limit for summarization by the model."
38
39    final_summary = " ".join(summaries)
40    return final_summary
41
42def summarize_pdf(text, model, tokenizer, device):
43    sentences = sent_tokenize(text)
44    chunks = [sentences[i:i+10] for i in range(0, len(sentences), 10)]
45    summaries = []
46
47    for chunk in chunks:
48        chunk_text = " ".join(chunk)
49        inputs = tokenizer(chunk_text, return_tensors="pt", max_length=1024, truncation=True)
50        input_ids = inputs["input_ids"].to(device)
51        attention_mask = inputs["attention_mask"].to(device)
52
53        summary_ids = model.generate(input_ids=input_ids, attention_mask=attention_mask, length_penalty=0.8, num_beams=8, max_length=128)
54        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
55        summaries.append(summary)
56
57    final_summary = " ".join(summaries)
58    return final_summary
59
60def read_pdf(uploaded_file):
61    reader = PdfReader(uploaded_file)
62    text = ""
63    for page in reader.pages:
64        text += page.extract_text()
65    return text
66
67def extractive_summary(text):
68    if len(text) > 2048:
69        raise IndexError("Input size exceeds the maximum limit for summarization by the model.")
70    summary = summarizer(text, max_length=500, min_length=80, do_sample=False)
71    return summary[0]['summary_text']
72
73def abstractive_summary(text):
74    if len(text) > 2048:
75        raise IndexError("Input size exceeds the maximum limit for summarization by the model.")
76    summary = summarizer(text, max_length=500, min_length=80, do_sample=True)
77    return summary[0]['summary_text']
78