README.md

metadata

library_name: transformers

Model Card for X-LoRA-Gemma-7b

X-LoRA-Gemma combines protein, chemical, bio-inspired and mechanics of materials capabilities. We use a set of four LoRA adapters, defined as follows:

1. Bioinspired materials
2. Mechanics and materials
3. Protein mechanics tasks (featuring generative sequence-to-property and inverse capabilities)
4. Quantum-mechanics based molecular properties QM9 (featuring generative SMILES-to-property and inverse capabilities

The model has a variety of capabilities, including designing proteins, designing molecules, and property calculations.

import torch
from xlora.xlora_utils import load_model  

XLoRa_model_name = 'lamm-mit/x-lora-gemma-7b'

model,tokenizer=load_model(model_name = XLoRa_model_name, 
                           device='cuda:0',
                           use_flash_attention_2=True, 
                           dtype=torch.bfloat16,
                            )
eos_token_id= tokenizer('<end_of_turn>', add_special_tokens=False, ) ['input_ids'][0]

def generate_XLoRA_Gemma (system_prompt='You a helpful assistant. You are familiar with materials science. ',
                     prompt='What is spider silk in the context of bioinspired materials?',
                     repetition_penalty=1.,num_beams=1,num_return_sequences=1,
                     top_p=0.9, top_k=256, temperature=.5,max_new_tokens=512, verbatim=False, eos_token=None, 
                     add_special_tokens=True, prepend_response='', 
                         ):
    if eos_token==None:
        eos_token= tokenizer.eos_token_id

    if system_prompt==None:
        messages=[ {"role": "user", "content": prompt},  ]
    else:
        messages=[ {"role": "user", "content": system_prompt+prompt},  ]
    txt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, )
    txt=txt+prepend_response
     
    inputs = tokenizer(txt, add_special_tokens  =add_special_tokens, return_tensors ='pt').to(device)
    with torch.no_grad():
         
          outputs = model.generate(input_ids = inputs["input_ids"],
                                   attention_mask = inputs["attention_mask"] , # This is usually done automatically by the tokenizer
                                   max_new_tokens=max_new_tokens,
                                   temperature=temperature, #value used to modulate the next token probabilities.
                                   num_beams=num_beams,
                                   top_k = top_k,
                                   top_p = top_p,
                                   num_return_sequences = num_return_sequences,
                                   eos_token_id=eos_token,
                                   pad_token_id = eos_token,
                                   do_sample =True,#skip_prompt=True,
                                   repetition_penalty=repetition_penalty,
                                   )
    return tokenizer.batch_decode(outputs[:,inputs["input_ids"].shape[1]:].detach().cpu().numpy(), skip_special_tokens=True)
    

Then, use as follows:

from IPython.display import display, Markdown
q='''What is graphene?'''
res=generate_XLoRA_Gemma( system_prompt='You design materials.', prompt=q, max_new_tokens=1024, temperature=0.3, eos_token=eos_token_id)
display (Markdown(res))

Example: Molecular design

def design_from_target(
    model,
    tokenizer,
    target,
    temperature=0.1,
    num_beams=1,
    top_k=50,
    top_p=0.95,
    repetition_penalty=1.0,
    messages=[]
):
    # Format the target line for molecular property generation
    line = f'GenerateMolecularProperties<{return_str(target)}>'
    
    # Add the line to the message history
    messages.append({"role": "user", "content": line})
    
    # Apply chat template with optional tokenization
    line = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # Generate response with specified parameters
    result = generate_response(
        model,
        tokenizer,
        text_input=line,
        num_return_sequences=1,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        max_new_tokens=256
    )[0]
    
    return result

Use case:

import numpy as np
target = np.random.rand(12)
SMILES=design_from_target (model, tokenizer, target, messages=[]])
print (SMILES)

Calculate molecular properties:

def properties_from_SMILES(
    model,
    tokenizer,
    target,
    temperature=0.1,
    top_k=128,
    top_p=0.9,
    num_beams=1,
    repetition_penalty=1.0
):
    # Format the target line for molecular property calculation
    line = f'CalculateMolecularProperties<{target}>'
    
    # Initialize messages and add the formatted line
    messages = [{"role": "user", "content": line}]
    
    # Apply chat template with optional tokenization
    line = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # Generate response with specified parameters
    result = generate_response(
        model,
        tokenizer,
        text_input=line,
        num_return_sequences=1,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        max_new_tokens=256
    )[0]
    
    # Extract relevant part of the result and convert to float list
    result = extract_start_and_end(result, start_token='[', end_token=']')
    return [float(i) for i in result.split(',')]