z-cao/DrugSpace-mntp-8B

DrugSpace MNTP 8B

DrugSpace-mntp-8B is the general-release Stage I model of DrugSpace. It was initialized from meta-llama/Meta-Llama-3.1-8B-Instruct and adapted to biomedical literature using the Masked Next Token Prediction (MNTP) objective with bidirectional attention.

This repository contains a complete merged model. A separate MNTP adapter is not required.

Model family

The general DrugSpace pipeline consists of:

• Stage I: z-cao/DrugSpace-mntp-8B
• Stage II: z-cao/DrugSpace-openfda-lora

For evaluation-oriented experiments using the pre-2020 Stage I corpus, use:

• z-cao/DrugSpace-mntp-eval-8B
• z-cao/DrugSpace-openfda-lora-eval

Usage

The model should be loaded through LLM2Vec with bidirectional attention enabled.

import torch
from llm2vec import LLM2Vec

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

model = LLM2Vec.from_pretrained(
    base_model_name_or_path="z-cao/DrugSpace-mntp-8B",
    peft_model_name_or_path="z-cao/DrugSpace-openfda-lora",
    enable_bidirectional=True,
    pooling_mode="mean",
    max_length=512,
    doc_max_length=400,
    skip_instruction=True,
    device_map=device,
    torch_dtype=dtype,
)

texts = [
    "A complete English description of a drug.",
    "Another drug description for comparison.",
]

embeddings = model.encode(texts, show_progress_bar=True)
print(embeddings.shape)

To use only the Stage I MNTP model, omit peft_model_name_or_path.

Training data

The model was trained on the full PubTator-derived Stage I biomedical literature corpus.

The corpus contains biomedical literature text associated with drug mentions and was prepared for continued representation learning from the Llama 3.1 8B Instruct initialization.

Training procedure

• Objective: Masked Next Token Prediction
• Attention: Bidirectional
• Training epochs: 1
• Maximum sequence length: 512
• Per-device training batch size: 32
• Masking probability: 0.20
• LoRA rank during MNTP training: 16
• Warm-up steps: 1,000
• Precision: bfloat16
• Attention implementation: FlashAttention 2
• Gradient checkpointing: Enabled

The MNTP LoRA weights were merged into the Llama base model to produce the complete checkpoint distributed in this repository.

Intended use

This model is intended for:

• use as the Stage I base model for z-cao/DrugSpace-openfda-lora;
• biomedical and drug-description embedding;
• semantic similarity;
• retrieval and ranking;
• continued research on biomedical representation learning.

Inputs should be complete English biomedical or drug descriptions. Consistent writing style is recommended when comparing embedding similarities.

Limitations

• This model is intended for representation learning rather than text generation or conversational use.
• Biomedical literature may contain reporting, coverage, and publication biases.
• Embedding similarity does not establish therapeutic equivalence, safety, efficacy, or causal relationships.
• This model is for research use and must not replace professional medical judgment.

License

This model is derived from meta-llama/Meta-Llama-3.1-8B-Instruct. Use is subject to the applicable Llama 3.1 license and acceptable-use terms.