news_summarizer / summerizer.py
Kota Takahashi
説明文を修正
fbd2a16
raw
history blame contribute delete
No virus
1.4 kB
from transformers import pipeline, T5Tokenizer, T5ForConditionalGeneration
class TextSummarizer:
model_name = "sonoisa/t5-base-japanese"
tokenizer_name = "sonoisa/t5-base-japanese"
def __init__(self):
"""
TextSummarizerクラスを初期化し、トークナイザ、モデル、パイプラインを設定
Parameters:
- なし
Returns:
- なし
"""
# トークナイザを個別に初期化し、legacy=Falseを指定
self.tokenizer = T5Tokenizer.from_pretrained(self.tokenizer_name, legacy=False)
# モデルを個別に初期化
self.model = T5ForConditionalGeneration.from_pretrained(self.model_name)
# パイプラインを初期化
self.summarizer = pipeline("summarization", model=self.model, tokenizer=self.tokenizer)
def summarize(self, text, max_length=20, min_length=10):
"""
テキストを要約
Parameters:
- text (str): 要約する対象のテキスト。
- max_length (int): 要約の最大長 (デフォルトは20)
- min_length (int): 要約の最小長 (デフォルトは10)
Returns:
- summary_text (str): 要約されたテキスト
"""
summary = self.summarizer(text, max_length=max_length, min_length=min_length, do_sample=False)
return summary[0]['summary_text']