SummaryProject / src /inference_t5.py
EstelleSkwarto's picture
ajout modele T5 dans API
ef05d9e
raw
history blame
1.85 kB
"""
Allows to predict the summary for a given entry text
"""
import torch
import nltk
import contractions
import re
import string
nltk.download('stopwords')
nltk.download('punkt')
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
def clean_data(texts):
texts = texts.lower()
texts = contractions.fix(texts)
texts = texts.translate(str.maketrans("", "", string.punctuation))
texts = re.sub(r'\n',' ',texts)
return texts
def inferenceAPI_t5(text: str) -> str:
"""
Predict the summary for an input text
--------
Parameter
text: str
the text to sumarize
Return
str
The summary for the input text
"""
# definition des parametres d'entree pour le modèle
text = clean_data(text)
device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
tokenizer= (AutoTokenizer.from_pretrained("./summarization_t5"))
# chargement du modele local
model = (AutoModelForSeq2SeqLM
.from_pretrained("./summarization_t5")
.to(device))
text_encoding = tokenizer(
text,
max_length=1024,
padding='max_length',
truncation=True,
return_attention_mask=True,
add_special_tokens=True,
return_tensors='pt'
)
generated_ids = model.generate(
input_ids=text_encoding['input_ids'],
attention_mask=text_encoding['attention_mask'],
max_length=128,
num_beams=8,
length_penalty=0.8,
early_stopping=True
)
preds = [
tokenizer.decode(gen_id, skip_special_tokens=True, clean_up_tokenization_spaces=True)
for gen_id in generated_ids
]
return "".join(preds)
if __name__ == "__main__":
text = input('Entrez votre phrase à résumer : ')
print('summary:',inferenceAPI(text))