Views
No views yet
1git clone https://huggingface.co/gsri-18/lora_finetuned_bert_on_body_parts_ner_dataset_synthetic
2cd lora_finetuned_bert_on_body_parts_ner_dataset_syntheticpip install transformers peft torch prettytable1import torch
2from transformers import BertTokenizerFast, BertForTokenClassification
3from peft import PeftModel
4from prettytable import PrettyTable # For tabular output
5
6# Device configuration
7device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8
9# Hugging Face model repository name
10base_model_name = "bert-base-uncased"
11lora_model_name = "gsri-18/lora_finetuned_bert_on_body_parts_ner_dataset_synthetic"
12
13# Load the tokenizer
14tokenizer = BertTokenizerFast.from_pretrained(base_model_name)
15
16# Load the base model
17base_model = BertForTokenClassification.from_pretrained(base_model_name, num_labels=3).to(device)
18
19# Load the LoRA fine-tuned model
20model = PeftModel.from_pretrained(base_model, lora_model_name).to(device)
21
22# Set the model to evaluation mode
23model.eval()
24
25# Function to predict entities with confidence scores
26def predict_entities(sentence, label_mapping):
27 # Tokenize the sentence
28 inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True).to(device)
29 with torch.no_grad():
30 # Get model outputs
31 outputs = model(**inputs)
32
33 # Process logits
34 logits = outputs.logits
35 probabilities = torch.nn.functional.softmax(logits, dim=-1)
36 predictions = torch.argmax(logits, dim=-1).squeeze().cpu().numpy()
37 confidence_scores = torch.max(probabilities, dim=-1).values.squeeze().cpu().numpy()
38 tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze().cpu().numpy())
39
40 # Extract entities
41 entities = []
42 for i, (token, prediction, confidence) in enumerate(zip(tokens, predictions, confidence_scores)):
43 if token.startswith("##"): # Skip subword tokens
44 continue
45 label = label_mapping[prediction]
46 if label != "O": # Skip tokens labeled as "O"
47 start = sentence.find(token)
48 end = start + len(token) - 1
49 entities.append({
50 "token": token,
51 "label": label,
52 "start": start,
53 "end": end,
54 "confidence": round(float(confidence), 8) # Round to 8 decimal places
55 })
56
57 return entities
58
59# Process and print predictions for multiple sentences
60def process_and_print_predictions(sentences, label_mapping):
61 for i, sentence in enumerate(sentences, 1):
62 print(f"\nSentence {i}: {sentence}")
63 entities = predict_entities(sentence, label_mapping)
64
65 if not entities:
66 print("No body parts were detected in the sentence.")
67 else:
68 # Create a table for better formatting
69 table = PrettyTable()
70 table.field_names = ["Token", "Label", "Start", "End", "Confidence"]
71
72 for entity in entities:
73 table.add_row([entity['token'], entity['label'], entity['start'], entity['end'], f"{entity['confidence']:.8f}"])
74
75 print(table)
76
77# Test sentences
78test_sentences = [
79 "The arm connects to the shoulder and elbow.",
80 "He felt pain in his knee and spinal cord after running.",
81 "The Named Entity Recognition model is working well."
82]
83
84label_mapping = {0: "O", 1: "B-BODY", 2: "I-BODY"}
85
86# Display results for all sentences
87process_and_print_predictions(test_sentences, label_mapping)+----------+--------+-------+-----+------------+
| Token | Label | Start | End | Confidence |
+----------+--------+-------+-----+------------+
| arm | B-BODY | 4 | 6 | 0.99613851 |
| shoulder | B-BODY | 24 | 31 | 0.99928027 |
| elbow | B-BODY | 37 | 41 | 0.99879104 |
+----------+--------+-------+-----+------------++--------+--------+-------+-----+------------+
| Token | Label | Start | End | Confidence |
+--------+--------+-------+-----+------------+
| knee | B-BODY | 20 | 23 | 0.99953437 |
| spinal | B-BODY | 29 | 34 | 0.99941099 |
| cord | I-BODY | 36 | 39 | 0.99498111 |
+--------+--------+-------+-----+------------+adapter_config.json: Configuration for the LoRA adapter.adapter_model.safetensors: The fine-tuned LoRA adapter weights.tokenizer.json, vocab.txt, tokenizer_config.json: Tokenizer files.model-bert-synthetic-bpr-lora-large.pth: The base model weights.