Views
No views yet
1import os
2import torch
3import torch.nn.functional as F
4from transformers import AutoTokenizer, BertForTokenClassification
5from tqdm import tqdm
6
7model_path = MatDetector_ckp
8# you can download matbert at https://github.com/lbnlp/MatBERT
9tokenizer_path = '/matbert-base-cased'
10input_file = 'TARGET.txt'
11output_directory = './'
12
13
14tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=False, do_lower_case=False)
15model = BertForTokenClassification.from_pretrained(model_path).half()
16
17device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
18model.to(device)
19model.eval()
20
21label_map = {0: "O", 1: "B-matname", 2: "I-matname", 3: "B-mf", 4: "I-mf"}
22
23def process_single_word(word, tokenizer, model, device):
24 tokenized = tokenizer(word, return_tensors="pt", truncation=True, max_length=128)
25 input_ids = tokenized["input_ids"].to(device)
26 attention_mask = tokenized["attention_mask"].to(device)
27
28 with torch.no_grad():
29 outputs = model(input_ids=input_ids, attention_mask=attention_mask)
30 logits = outputs.logits
31 probabilities = F.softmax(logits, dim=2) # (batch_size=1, seq_len, num_labels)
32
33 return tokenized, probabilities
34
35
36def determine_label(tokenized, probabilities, label_map):
37 tokens = tokenizer.convert_ids_to_tokens(tokenized["input_ids"][0].tolist())
38 probs = probabilities[0] # (seq_len, num_labels)
39
40 token_labels = []
41 for token, prob in zip(tokens, probs):
42 if token in ["[CLS]", "[SEP]", "[PAD]"]:
43 continue
44
45 clean_token = token[2:] if token.startswith("##") else token
46 max_label = prob.argmax().item()
47 label_name = label_map[max_label]
48
49 token_labels.append(label_name)
50
51 label_counts = {}
52 for label in token_labels:
53 if label not in label_counts:
54 label_counts[label] = 0
55 label_counts[label] += 1
56
57 final_label = max(label_counts, key=label_counts.get) if label_counts else "O"
58
59 return final_label
60
61
62with open(os.path.join(output_directory, "mf.txt"), "w") as mf_file, \
63 open(os.path.join(output_directory, "matname.txt"), "w") as matname_file, \
64 open(os.path.join(output_directory, "o_tags.txt"), "w") as o_file:
65
66 with open(input_file, 'r') as file:
67 lines = [line.strip() for line in file.readlines() if line.strip()]
68 total_lines = len(lines)
69
70 with tqdm(total=total_lines, desc="Processing words", unit="words") as progress_bar:
71 for original_word in lines:
72 tokenized, probabilities = process_single_word(original_word, tokenizer, model, device)
73 final_label = determine_label(tokenized, probabilities, label_map)
74
75 if final_label == "O":
76 o_file.write(f"{original_word}\n")
77 elif final_label in ["B-mf", "I-mf"]:
78 mf_file.write(f"{original_word}\n")
79 elif final_label in ["B-matname", "I-matname"]:
80 matname_file.write(f"{original_word}\n")
81
82 progress_bar.update(1)
83
84print("Processing completed. Files saved as mf.txt, matname.txt, and o_tags.txt.")
85