Update README.md

README.md

@@ -69,7 +69,9 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 model_id = "protgpt3/ProtGPT3-1OB"  # Replace with the final checkpoint name
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+# Load tokenizer for generation
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True,add_bos_token=True, add_eos_token=False)
+
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype=torch.bfloat16,
@@ -80,12 +82,12 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 ```
 
-Generate a protein sequence:
+### Generate a protein sequence
 
 ```python
 import torch
 
-prompt = ""  # Optionally provide an amino-acid prefix or model-specific direction token
+prompt = ""  # Optionally provide an amino-acid prefix or model-specific direction
 
 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
@@ -97,17 +99,20 @@ with torch.no_grad():
         temperature=0.8,
         top_p=0.9,
         eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
     )
 
 sequence = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-print(sequence)
+print(sequence) # output includes directional token "1" or "2" to denote if sequence was generated N-to-C or C-to-N
 ```
 
-Generate from an amino-acid prefix:
+### Generate from an amino-acid prefix
 
 ```python
-prefix = "MKT"
+import torch
+
+# forward N-to-C generation with special token "1" 
+prefix = "1MKT" # use special token "2" instead of "1" for reverse  C-to-N generation
 
 inputs = tokenizer(prefix, return_tensors="pt").to(model.device)
 
@@ -126,6 +131,40 @@ sequence = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 print(sequence)
 ```
 
+### Batch generation
+
+```python
+import torch
+
+prompts = [
+    "",
+    "1MKT", # N-to-C generation
+    "2MAV", # C-to-N generation
+]
+
+inputs = tokenizer(
+    prompts,
+    return_tensors="pt",
+    padding=True,
+).to(model.device)
+
+with torch.no_grad():
+    output_ids = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.8,
+        top_p=0.9,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.bos_token_id,
+    )
+
+sequences = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+
+for sequence in sequences:
+    print(sequence)
+```
+
 ## Training Details
 
 ### Training Data