Spaces:
Sleeping
Sleeping
File size: 6,132 Bytes
1361c8e b75daf4 1361c8e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel
import numpy as np
from sklearn.preprocessing import StandardScaler
from joblib import load
import random
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
# set_seed(42)
class BertEmbedder:
def __init__(self, model_name = "./embedding-model/matscibert", device="cuda"):
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModel.from_pretrained(model_name).to(device)
self.device = device
self.model.eval()
@torch.no_grad()
def encode(self, text):
inputs = self.tokenizer(
text,
padding=True,
truncation=True,
max_length=512,
return_tensors="pt"
).to(self.device)
outputs = self.model(**inputs)
last_hidden = outputs.last_hidden_state
mask = inputs.attention_mask.unsqueeze(-1)
embedding = (last_hidden * mask).sum(1) / mask.sum(1)
return embedding.cpu().numpy().squeeze().tolist()
class FeatureExtractor(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(779, 1024),
nn.ReLU(),
nn.Linear(1024, 512),
nn.ReLU(),
nn.Linear(512, 256),
nn.Linear(256, 32),
)
def forward(self, x):
return self.net(x).squeeze(-1)
class MCDropoutModel(nn.Module):
def __init__(self, input_dim=32, dropout_rate=0.1):
super().__init__()
# self.fc1 = nn.Linear(input_dim, 64)
self.fc2 = nn.Linear(32, 128)
self.fc3 = nn.Linear(128, 32)
self.fc4 = nn.Linear(32, 1)
self.dropout = nn.Dropout(dropout_rate)
def forward(self, x):
# x = F.leaky_relu(self.fc1(x), negative_slope=0.05)
x = F.leaky_relu(self.dropout(self.fc2(x)), negative_slope=0.05)
x = F.leaky_relu(self.dropout(self.fc3(x)), negative_slope=0.05)
x = self.fc4(x)
return x.squeeze(-1)
def mc_dropout_predict(model, X, n_samples = 50):
set_seed(42)
device = next(model.parameters()).device
X_tensor = torch.FloatTensor(X).to(device)
model.eval()
for m in model.modules():
if isinstance(m, nn.Dropout):
m.train()
predictions = []
with torch.no_grad():
for _ in range(n_samples):
pred = model(X_tensor)
predictions.append(pred.cpu().numpy())
predictions = np.array(predictions)
return predictions.mean(axis=0), predictions.std(axis=0)
def predict(Al_content=None,
Nb_content=None,
Ta_content=None,
Ti_content=None,
Zr_content=None,
Mo_content=None,
V_content=None,
Cr_content=None,
W_content=None,
Hf_content=None,
Ni_content=None,
process_text = None,
scaler_model = "./ANN-modeling/mc_dropout_scaler.pkl",
model_path = './ANN-modeling/mc_dropout_best_vam_only.pth',
embedding_model = "./embedding-model/matscibert",
extract_model = "./DANN-modeling/feature_extractor.pth",
device = 'cuda'):
comp_list = [
Al_content or 0.0,
Nb_content or 0.0,
Ta_content or 0.0,
Ti_content or 0.0,
Zr_content or 0.0,
Mo_content or 0.0,
V_content or 0.0,
Cr_content or 0.0,
W_content or 0.0,
Hf_content or 0.0,
Ni_content or 0.0,
]
comp_list = [float(x) for x in comp_list]
if process_text is not None:
# print("Embedding process text...")
embedder = BertEmbedder(embedding_model)
emb = embedder.encode(process_text)
proc_emb = np.array(emb)
# print("Extracting features...")
extractor = FeatureExtractor().to(device)
extractor.load_state_dict(torch.load(extract_model, map_location=device))
# combinate process and composition
# X = np.hstack([comp_list, proc_emb])
comp_array = np.atleast_2d(comp_list) # Shape: (n_samples, 11)
proc_emb_array = np.tile(proc_emb, (comp_array.shape[0], 1)) # Repeat proc_emb for each sample
X = np.hstack([comp_array, proc_emb_array])
lantent_features = extractor(torch.FloatTensor(X).to(device)).cpu().detach().numpy()
# print("Predicting...")
model = MCDropoutModel().to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
# scaler = StandardScaler()
scaler = load(scaler_model)
lantent_features_scaled = scaler.transform(lantent_features)
pred, std = mc_dropout_predict(model, lantent_features_scaled)
return np.round(pred, 2), np.round(std, 2)
if __name__ == "__main__":
# set_seed(42)
process_text = "prepared by laser powder bed fusion technique using a feedstock of pre-alloyed powders with optimized process parameters: laser power 230 W, scanning speed 900 mm/s, hatch spacing 62 μm, and layer thickness 30 μm."
ys, std = predict(Al_content=25,
Nb_content=25,
Ta_content=None,
Ti_content=25,
Zr_content=None,
Mo_content=None,
V_content=25,
Cr_content=None,
W_content=None,
Hf_content=None,
Ni_content=None,
process_text = process_text,
scaler_model = "./ANN-modeling/mc_dropout_scaler.pkl",
model_path = './ANN-modeling/mc_dropout_best_vam_only.pth',
embedding_model = "./embedding-model/matscibert",
extract_model = "./DANN-modeling/feature_extractor.pth",
device = 'cuda')
print(ys, std) |