Views
No views yet
pip install transformers torch evaluate1from transformers import BertTokenizerFast, BertForSequenceClassification
2import torch
3
4# Load the fine-tuned model and tokenizer
5model_path = "final_relation_extraction_model"
6tokenizer = BertTokenizerFast.from_pretrained(model_path)
7model = BertForSequenceClassification.from_pretrained(model_path)
8model.eval()
9
10# Example input with entity markers
11text = "笔名：[SUBJ] 木斧 [/SUBJ] 出生地：[OBJ] 成都 [/OBJ]"
12
13# Tokenize input
14inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
15
16# Inference
17with torch.no_grad():
18 outputs = model(**inputs)
19 logits = outputs.logits
20 predicted_class = torch.argmax(logits, dim=1).item()
21
22# Map prediction to relation label
23label_mapping = {0: "出生地", 1: "出生日期", 2: "民族", 3: "职业"} # Customize based on your labels
24predicted_relation = label_mapping[predicted_class]
25print(f"Predicted Relation: {predicted_relation}")[SUBJ] and [OBJ].
├── final_relation_extraction_model/
│ ├── config.json
│ ├── pytorch_model.bin # Fine-tuned Model
│ ├── tokenizer_config.json
│ ├── special_tokens_map.json
│ ├── tokenizer.json
│ ├── vocab.txt
│ └── added_tokens.json
├── relationship-extraction.ipynb # Training notebook
└── README.md # Model documentation[SUBJ] entity_name [/SUBJ][OBJ] entity_name [/OBJ]Input: "笔名：[SUBJ] 木斧 [/SUBJ]原名：杨莆曾民族： [OBJ] 回族 [/OBJ]"
Output: "民族" (ethnicity relation)