Views
No views yet
chemistry drug-discovery generative-model smiles lstm rdkitPAD, BOS, EOS).1import torch
2from huggingface_hub import hf_hub_download
3from model import SmilesLSTM
4from tokenizer import CharVocab
5
6# Download model file (.pth) from Hugging Face
7model_path = hf_hub_download(
8 repo_id="GioFilo93/Molecule-Generation_LSTM-Based",
9 filename="lstm_molecule-generation.pth"
10)
11
12# Load tokenizer (must be the same vocab used in training)
13# Example: vocab.json or hardcoded CharVocab
14vocab = CharVocab.load("vocab.json")
15
16# Init model
17model = SmilesLSTM(
18 vocab_size=len(vocab), emb_dim=128,
19 hidden_size=512, num_layers=2, dropout=0.3
20)
21state = torch.load(model_path, map_location="cpu")
22model.load_state_dict(state)
23model.eval()
24
25# Sampling
26@torch.no_grad()
27def sample(n=5, max_len=120, temperature=0.9):
28 out = []
29 for _ in range(n):
30 token = torch.tensor([vocab.bos_id])
31 seq, h = [], None
32 for _ in range(max_len):
33 logits, h = model.step(token, h)
34 probs = torch.softmax(logits.squeeze(0) / temperature, dim=-1)
35 token = torch.multinomial(probs, 1)
36 if token.item() == vocab.eos_id:
37 break
38 seq.append(token.item())
39 out.append(vocab.decode(seq))
40 return out
41
42print(sample(n=5, temperature=0.8))
43