Spaces:
Runtime error
Runtime error
""" | |
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)) | |