Views
No views yet
1from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
2import re
3
4model_path = "cassandra-themis/lsg-ner-phrases-16384"
5
6model = AutoModelForTokenClassification.from_pretrained(model_path, trust_remote_code=True, use_auth_token=True)
7tokenizer = AutoTokenizer.from_pretrained(model_path, use_auth_token=True)
8ner_pipe = pipeline("token-classification", model=model, tokenizer=tokenizer)
9
10
11document = "My document"
12document_flattened = re.sub(r'(\s|\t|\n)+', r' ', document).strip()
13
14prediction = ner_pipe(document_flattened, aggregation_strategy="simple")
15
16sentences = []
17for i in range(len(prediction) - 1):
18 sentences.append(document_flattened[prediction[i]["start"]:prediction[i+1]["start"]].strip())
19print("\n".join(sentences))