File size: 3,621 Bytes
9b54ac6
 
 
 
 
 
 
 
 
 
 
 
979bf1b
 
 
 
 
9b54ac6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
---
license: apache-2.0
base_model:
- meta-llama/Llama-3.3-70B-Instruct
- unsloth/Meta-Llama-3.1-8B-Instruct
tags:
- biology
- chemistry
---

# Pro-1-preview

[![GitHub](https://img.shields.io/badge/GitHub-michaelhla/pro--1-181717?logo=github)](https://github.com/michaelhla/pro-1)
[![Twitter](https://img.shields.io/badge/Twitter-@hla__michael-1DA1F2?logo=twitter&style=social)](https://twitter.com/hla_michael)
[![Hugging Face](https://img.shields.io/badge/🤗%20Hugging%20Face-mhla/pro--1-yellow)](https://huggingface.co/mhla/pro-1)
[![Blog Post](https://img.shields.io/badge/Blog-pro--1-red)](https://michaelhla.com/blog/pro1)

Pro-1 is a reasoning model trained using GRPO towards a physics based reward function for protein stability.

It takes in a protein sequence + text description of the protein + effects of previous engineering attempts, reasons over the information given, and proposes modifications to improve the stability of the given sequence. 


## LORA checkpoints 
| Model | Checkpoint |
|------------|-------------|
| 8b base GRPO | [best-checkpoint](https://huggingface.co/mhla/pro-1/tree/main/best-checkpoint) |
| 8b creative reward | [creativity-lm-grpo-mega-run-full](https://huggingface.co/mhla/pro-1/tree/main/creativity-lm-grpo-mega-run-full) |
| 8b creative + specificity reward (default) | [all-lm-grpo-mega-run](https://huggingface.co/mhla/pro-1/tree/main/all-lm-grpo-mega-run-full) |
| 70b SFT only | [llama_70b_4bit_sft_lora_model](https://huggingface.co/mhla/pro-1/tree/main/llama_70b_4bit_sft_lora_model) |


## Example Usage

```python
from unsloth import FastLanguageModel
from transformers import TextIteratorStreamer
import threading

def run_protein_engineering_example():
    # Load the model and tokenizer
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name="unsloth/meta-Llama-3.1-8B-Instruct",
        max_seq_length=32768,
        load_in_4bit=True,
        fast_inference=True,
        max_lora_rank=32,
        gpu_memory_utilization=0.6,
    )
    
    # Load the protein engineering adapter weights
    model.load_adapter("your-username/protein-engineering-llama-3.1")
    FastLanguageModel.for_inference(model)
    
    protein_sequence = "MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKPLSVSYDQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTY"
    
    prompt = f"""

...{STRUCTURED PROMPT SEE https://github.com/michaelhla/pro-1 FOR CORRECT USAGE}...

"""

    # Initialize the streamer for text generation
    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
    
    # Set up generation parameters
    generation_kwargs = dict(
        input_ids=tokenizer(prompt, return_tensors="pt").input_ids.to(model.device),
        streamer=streamer,
        max_new_tokens=4096,
        temperature=0.9,
        top_p=0.95,
        do_sample=True
    )
    
    # Create a thread to run the generation
    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    # Print the response as it streams
    print("Model response (streaming):")
    for new_text in streamer:
        print(new_text, end="", flush=True)
    
    thread.join()  # Ensure generation is complete

if __name__ == "__main__":
    run_protein_engineering_example()

```


Note: While the model was specifically trained on enzymes, it should work for any protein sequence. Curious to hear if this is true!

Disclaimer: This is a preview version and as a result the model can be very dumb. Always double check sure your modified sequences have the correct mutations applied. Assume all references from the model are hallucinated.