Views
No views yet
1pip install -U transformers peft accelerate bitsandbytes
2pip install git+https://github.com/EricLBuehler/xlora.git
3pip install -U rdkit scikit-learn tqdm pandas{guidance} for inference:pip install guidance
1import torch
2from xlora.xlora_utils import load_model
3
4XLoRa_model_name = 'lamm-mit/x-lora-gemma-7b'
5
6model,tokenizer=load_model(model_name = XLoRa_model_name,
7 device='cuda:0',
8 use_flash_attention_2=True,
9 dtype=torch.bfloat16,
10 )
11eos_token_id= tokenizer('<end_of_turn>', add_special_tokens=False, ) ['input_ids'][0]1def generate_XLoRA_Gemma (system_prompt='You a helpful assistant. You are familiar with materials science. ',
2 prompt='What is spider silk in the context of bioinspired materials?',
3 repetition_penalty=1.,num_beams=1,num_return_sequences=1,
4 top_p=0.9, top_k=256, temperature=.5,max_new_tokens=512, verbatim=False, eos_token=None,
5 add_special_tokens=True, prepend_response='',
6 ):
7 if eos_token==None:
8 eos_token= tokenizer.eos_token_id
9
10 if system_prompt==None:
11 messages=[ {"role": "user", "content": prompt}, ]
12 else:
13 messages=[ {"role": "user", "content": system_prompt+prompt}, ]
14 txt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, )
15 txt=txt+prepend_response
16
17 inputs = tokenizer(txt, add_special_tokens =add_special_tokens, return_tensors ='pt').to(device)
18 with torch.no_grad():
19
20 outputs = model.generate(input_ids = inputs["input_ids"],
21 attention_mask = inputs["attention_mask"] , # This is usually done automatically by the tokenizer
22 max_new_tokens=max_new_tokens,
23 temperature=temperature, #value used to modulate the next token probabilities.
24 num_beams=num_beams,
25 top_k = top_k,
26 top_p = top_p,
27 num_return_sequences = num_return_sequences,
28 eos_token_id=eos_token,
29 pad_token_id = eos_token,
30 do_sample =True,#skip_prompt=True,
31 repetition_penalty=repetition_penalty,
32 )
33 return tokenizer.batch_decode(outputs[:,inputs["input_ids"].shape[1]:].detach().cpu().numpy(), skip_special_tokens=True)
34 1from IPython.display import display, Markdown
2q='''What is graphene?'''
3res=generate_XLoRA_Gemma( system_prompt='You design materials.', prompt=q, max_new_tokens=1024, temperature=0.3, eos_token=eos_token_id)
4display (Markdown(res))
1def design_from_target(
2 model,
3 tokenizer,
4 target,
5 temperature=0.1,
6 num_beams=1,
7 top_k=50,
8 top_p=0.95,
9 repetition_penalty=1.0,
10 messages=[]
11):
12 # Format the target line for molecular property generation
13 line = f'GenerateMolecularProperties<{return_str(target)}>'
14
15 # Add the line to the message history
16 messages.append({"role": "user", "content": line})
17
18 # Apply chat template with optional tokenization
19 line = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
20
21 # Generate response with specified parameters
22 result = generate_response(
23 model,
24 tokenizer,
25 text_input=line,
26 num_return_sequences=1,
27 temperature=temperature,
28 top_k=top_k,
29 top_p=top_p,
30 max_new_tokens=256
31 )[0]
32
33 return result1import numpy as np
2target = np.random.rand(12)
3SMILES=design_from_target (model, tokenizer, target, messages=[]])
4print (SMILES)1def properties_from_SMILES(
2 model,
3 tokenizer,
4 target,
5 temperature=0.1,
6 top_k=128,
7 top_p=0.9,
8 num_beams=1,
9 repetition_penalty=1.0
10):
11 # Format the target line for molecular property calculation
12 line = f'CalculateMolecularProperties<{target}>'
13
14 # Initialize messages and add the formatted line
15 messages = [{"role": "user", "content": line}]
16
17 # Apply chat template with optional tokenization
18 line = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
19
20 # Generate response with specified parameters
21 result = generate_response(
22 model,
23 tokenizer,
24 text_input=line,
25 num_return_sequences=1,
26 temperature=temperature,
27 top_k=top_k,
28 top_p=top_p,
29 max_new_tokens=256
30 )[0]
31
32 # Extract relevant part of the result and convert to float list
33 result = extract_start_and_end(result, start_token='[', end_token=']')
34 return [float(i) for i in result.split(',')]