POET / predict.py
qanastek's picture
First commit
55be9d0
from transformers import CamembertTokenizer, CamembertForTokenClassification, TokenClassificationPipeline
OUTPUT_PATH = './'
tokenizer = CamembertTokenizer.from_pretrained(OUTPUT_PATH)
model = CamembertForTokenClassification.from_pretrained(OUTPUT_PATH)
pos = TokenClassificationPipeline(model=model, tokenizer=tokenizer)
def make_prediction(sentence):
labels = [l['entity'] for l in pos(sentence)]
return list(zip(sentence.split(" "), labels))
res = make_prediction("George Washington est allé à Washington")