Update README.md

README.md

@@ -27,119 +27,100 @@ For more technical details, please refer to our paper [GENERator: A Long-Context
 
 
 ## How to use
-### Simple example1:  generation
+### Example 1: Sequence Generation
 
 ```python
 
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load the tokenizer and model.
-tokenizer = AutoTokenizer.from_pretrained("GenerTeam/GENERator-eukaryote-3b-base", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("GenerTeam/GENERator-eukaryote-3b-base")
-config = model.config
+model = AutoModelForCausalLM.from_pretrained(
+    "GenerTeam/GENERator-eukaryote-1.2b-base",
+    attn_implementation="flash_attention_2",
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+).cuda().eval()
 
-max_length = config.max_position_embeddings
+tokenizer = AutoTokenizer.from_pretrained(
+    "GenerTeam/GENERator-eukaryote-1.2b-base",
+    trust_remote_code=True,
+)
 
 # Define input sequences.
 sequences = [
-    "ATGAGGTGGCAAGAAATGGGCTAC",
-    "GAATTCCATGAGGCTATAGAATAATCTAAGAGAAAT"
+    "ATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCG",
+    "ACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGT"
 ]
 
-def left_padding(sequence, padding_char='A', multiple=6):
-    remainder = len(sequence) % multiple
-    if remainder != 0:
-        padding_length = multiple - remainder
-        return padding_char * padding_length + sequence
-    return sequence
-
-def left_truncation(sequence, multiple=6):
-    remainder = len(sequence) % multiple
-    if remainder != 0:
-        return sequence[remainder:]
-    return sequence
-
-# Apply left_padding to all sequences
-# padded_sequences = [left_padding(seq) for seq in sequences]
-
-# Apply left_truncation to all sequences
-truncated_sequences = [left_truncation(seq) for seq in sequences]
-
-# Process the sequences
-sequences = [tokenizer.bos_token + sequence for sequence in truncated_sequences]
+# Truncate each sequence to the nearest multiple of 6
+processed_sequences = ["<s>" + seq[len(seq)%6:] for seq in sequences]
 
 # Tokenize the sequences
-tokenizer.padding_side = "left"
 inputs = tokenizer(
-    sequences,
+    processed_sequences,
     add_special_tokens=False,
     return_tensors="pt",
     padding=True,
-    truncation=True,
-    max_length=max_length
-)
+    padding_side="left",
+).to("cuda")
 
 # Generate the sequences
 with torch.inference_mode():
-    outputs = model.generate(**inputs, max_new_tokens=32, temperature=0.00001, top_k=1)
+    outputs = model.generate(**inputs, max_new_tokens=32, do_sample=False)
 
 # Decode the generated sequences
 decoded_sequences = tokenizer.batch_decode(outputs, skip_special_tokens=True)
 
 # Print the decoded sequences
 print(decoded_sequences)
-
-# It is expected to observe non-sense decoded sequences (e.g., 'AAAAAA')
-# The input sequences are too short to provide sufficient context.
 ```
 
-### Simple example2: embedding
+### Example 2: Embedding Extraction
 
 ```python
 
-
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("GENERator-eukaryote-3b-base", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("GENERator-eukaryote-3b-base")
+model = AutoModelForCausalLM.from_pretrained(
+    "GenerTeam/GENERator-eukaryote-1.2b-base",
+    attn_implementation="flash_attention_2",
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+).cuda().eval()
 
-# Get model configuration
-config = model.config
-max_length = config.max_position_embeddings
+tokenizer = AutoTokenizer.from_pretrained(
+    "GenerTeam/GENERator-eukaryote-1.2b-base",
+    trust_remote_code=True,
+)
 
-# Define input sequences
+# Define input sequences.
 sequences = [
-    "ATGAGGTGGCAAGAAATGGGCTAC",
-    "GAATTCCATGAGGCTATAGAATAATCTAAGAGAAAT"
+    "ATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCG",
+    "ACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGT"
 ]
 
 # Truncate each sequence to the nearest multiple of 6
-processed_sequences = [tokenizer.bos_token + seq[:len(seq)//6*6] for seq in sequences]
+processed_sequences = ["<s>" + seq[len(seq)%6:] for seq in sequences]
 
-# Tokenization
-tokenizer.padding_side = "right"
+# Tokenize the sequences
 inputs = tokenizer(
     processed_sequences,
-    add_special_tokens=True,
+    add_special_tokens=False,
     return_tensors="pt",
     padding=True,
-    truncation=True,
-    max_length=max_length
-)
+    padding_side="right",
+).to("cuda")
 
-# Model Inference
 with torch.inference_mode():
     outputs = model(**inputs, output_hidden_states=True)
 
 hidden_states = outputs.hidden_states[-1]
 attention_mask = inputs["attention_mask"]
 
-# Option 1: Last token (EOS) embedding
+# Option 1: Last token embedding
 last_token_indices = attention_mask.sum(dim=1) - 1
-eos_embeddings = hidden_states[torch.arange(hidden_states.size(0)), last_token_indices, :]
+last_token_embeddings = hidden_states[torch.arange(hidden_states.size(0)), last_token_indices, :]
 
 # Option 2: Mean pooling over all tokens
 expanded_mask = attention_mask.unsqueeze(-1).expand(hidden_states.size()).to(torch.float32)
@@ -147,16 +128,13 @@ sum_embeddings = torch.sum(hidden_states * expanded_mask, dim=1)
 mean_embeddings = sum_embeddings / expanded_mask.sum(dim=1)
 
 # Output
-print("EOS (Last Token) Embeddings:", eos_embeddings)
+print("Last Token Embeddings:", last_token_embeddings)
 print("Mean Pooling Embeddings:", mean_embeddings)
 
 # ============================================================================
-# Additional notes:
-# - The preprocessing step ensures sequences are multiples of 6 for 6-mer tokenizer
-# - For causal LM, the last token embedding (EOS) is commonly used
-# - Mean pooling considers all tokens including BOS and content tokens
-# - The choice depends on your downstream task requirements
-# - Both methods handle variable sequence lengths via attention mask
+# The choice depends on your downstream task requirements
+# - Last token embeddings capture more localized gene-level information (e.g., strand, codon phase).
+# - Mean pooling embeddings capture species-level information.
 # ============================================================================
 
 ```