ayushsinha commited on
Commit
159d35d
·
verified ·
1 Parent(s): b8875f4

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -5
README.md CHANGED
@@ -42,22 +42,19 @@ label_list = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC", "B-MISC
42
 
43
  ```
44
  def predict_entities(text, model):
45
- # ✅ Tokenize input text
46
  tokens = tokenizer(text, return_tensors="pt", truncation=True)
47
  tokens = {key: val.to(device) for key, val in tokens.items()} # Move to CUDA
48
 
49
- # ✅ Run model inference
50
  with torch.no_grad():
51
  outputs = model(**tokens)
52
 
53
  logits = outputs.logits # Extract logits
54
  predictions = torch.argmax(logits, dim=2) # Get highest probability labels
55
 
56
- # ✅ Convert token IDs back to words
57
  tokens_list = tokenizer.convert_ids_to_tokens(tokens["input_ids"][0])
58
  predicted_labels = [label_list[pred] for pred in predictions[0].cpu().numpy()]
59
 
60
- # ✅ Group subword tokens into whole words
61
  final_tokens = []
62
  final_labels = []
63
  for token, label in zip(tokens_list, predicted_labels):
@@ -67,7 +64,6 @@ def predict_entities(text, model):
67
  final_tokens.append(token)
68
  final_labels.append(label)
69
 
70
- # ✅ Display results (ignore special tokens)
71
  for token, label in zip(final_tokens, final_labels):
72
  if token not in ["[CLS]", "[SEP]"]:
73
  print(f"{token}: {label}")
 
42
 
43
  ```
44
  def predict_entities(text, model):
45
+
46
  tokens = tokenizer(text, return_tensors="pt", truncation=True)
47
  tokens = {key: val.to(device) for key, val in tokens.items()} # Move to CUDA
48
 
 
49
  with torch.no_grad():
50
  outputs = model(**tokens)
51
 
52
  logits = outputs.logits # Extract logits
53
  predictions = torch.argmax(logits, dim=2) # Get highest probability labels
54
 
 
55
  tokens_list = tokenizer.convert_ids_to_tokens(tokens["input_ids"][0])
56
  predicted_labels = [label_list[pred] for pred in predictions[0].cpu().numpy()]
57
 
 
58
  final_tokens = []
59
  final_labels = []
60
  for token, label in zip(tokens_list, predicted_labels):
 
64
  final_tokens.append(token)
65
  final_labels.append(label)
66
 
 
67
  for token, label in zip(final_tokens, final_labels):
68
  if token not in ["[CLS]", "[SEP]"]:
69
  print(f"{token}: {label}")