Views
No views yet
1
2import torch
3from transformers import AutoModelForTokenClassification, AutoTokenizer
4
5
6# Model name from Hugging Face model hub
7model_name = "zekun-li/geolm-base-toponym-recognition"
8
9# Load tokenizer and model
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11model = AutoModelForTokenClassification.from_pretrained(model_name)
12
13# Example input sentence
14input_sentence = "Minneapolis, officially the City of Minneapolis, is a city in the state of Minnesota and the county seat of Hennepin County."
15
16# Tokenize input sentence
17tokens = tokenizer.encode(input_sentence, return_tensors="pt")
18
19# Pass tokens through the model
20outputs = model(tokens)
21
22# Retrieve predicted labels for each token
23predicted_labels = torch.argmax(outputs.logits, dim=2)
24
25predicted_labels = predicted_labels.detach().cpu().numpy()
26
27# Decode predicted labels
28predicted_labels = [model.config.id2label[label] for label in predicted_labels[0]]
29
30# Print predicted labels
31print(predicted_labels)
32# ['O', 'B-Topo', 'O', 'O', 'O', 'O', 'O', 'B-Topo', 'O', 'O', 'O', 'O', 'O', 'O',
33# 'O', 'O', 'B-Topo', 'O', 'O', 'O', 'O', 'O', 'B-Topo', 'I-Topo', 'I-Topo', 'O', 'O', 'O']