@@ -40,15 +40,23 @@ def parse_fasta_sequences(fasta_path: Path) -> List[str]:
4040 return sequences
4141
4242def sample_peptides_from_fasta (fasta_path : Path , length : int , count : int ) -> List [str ]:
43+ print (f"Parsing FASTA file: { fasta_path } )
4344 sequences = parse_fasta_sequences (fasta_path )
45+
46+ print (f"Extracting { length } { len (sequences )} )
4447 all_subseqs = set () # Use set to automatically collapse duplicates
45- for seq in sequences :
48+
49+ # Use progress bar for subsequence extraction
50+ for seq in tqdm (sequences , desc = "Processing proteins" , unit = "protein" ):
4651 if len (seq ) >= length :
4752 for i in range (len (seq ) - length + 1 ):
4853 all_subseqs .add (seq [i :i + length ])
54+
4955 if not all_subseqs :
5056 raise ValueError (f"No subsequences of length { length } { fasta_path } )
5157
58+ print (f"Found { len (all_subseqs )} { length } )
59+
5260 # Convert set back to list for sampling
5361 unique_subseqs = list (all_subseqs )
5462
@@ -58,6 +66,7 @@ def sample_peptides_from_fasta(fasta_path: Path, length: int, count: int) -> Lis
5866 return unique_subseqs
5967
6068 # Sample without replacement
69+ print (f"Sampling { count } )
6170 return random .sample (unique_subseqs , k = count )
6271
6372def generate_llm_peptides (length : int , count : int , model_name : str = "protgpt2" , top_k : int = 950 , top_p : float = 0.9 , repetition_penalty : float = 1.2 ) -> List [str ]:
@@ -332,9 +341,11 @@ def generate_fake_proteome(num_proteins: int, target_lengths: List[int], model_n
332341 sys .exit (1 )
333342
334343def write_fasta (peptides : List [str ], output_path : Path , prefix : str = "peptide" ):
344+ print (f"Writing { len (peptides )} { output_path } )
335345 with open (output_path , 'w' ) as f :
336- for i , pep in enumerate (peptides , 1 ):
337- f .write (f">{ prefix } { i } \n { pep } \n " )
346+ for i , pep in enumerate (tqdm (peptides , desc = "Writing peptides" , unit = "peptide" )):
347+ f .write (f">{ prefix } { i + 1 } \n { pep } \n " )
348+ print (f"✅ Successfully wrote { len (peptides )} { output_path } )
338349
339350def main ():
340351 parser = argparse .ArgumentParser (description = "Generate control peptides for neoantigen analysis." )
@@ -381,39 +392,50 @@ def main():
381392 sys .exit (1 )
382393 peptides = sample_peptides_from_fasta (args .fasta_file , args .length , args .count )
383394 elif args .source == 'llm' :
384- # Interactive workflow for LLM-based peptide generation
385- print (f"\n Generating { args .count } { args .length } )
386- print ("This approach generates a fake proteome first, then samples peptides from it." )
387-
388- # Check if user has existing proteome
389- has_existing = get_user_input ("\n Do you have an existing ProtGPT2-generated proteome file? (y/n): " ).lower ().startswith ('y' )
390-
391- if has_existing :
392- proteome_path = get_existing_proteome_path ()
393- print (f"Using existing proteome: { proteome_path } )
395+ # Choose workflow based on the LLM model
396+ if args .llm_model .lower () == 'protgpt2' :
397+ # Interactive proteome workflow for ProtGPT2 (to avoid M bias)
398+ print (f"\n Generating { args .count } { args .length } )
399+ print ("This approach generates a fake proteome first, then samples peptides from it." )
400+ elif args .llm_model .lower () == 'esm2' :
401+ # Direct generation for ESM2 (no M bias issue)
402+ print (f"\n Generating { args .count } { args .length } )
403+ peptides = generate_llm_peptides (args .length , args .count , args .llm_model , args .top_k , args .top_p , args .repetition_penalty )
394404 else :
395- # Ask if user wants to generate new proteome
396- generate_new = get_user_input ("Would you like to generate a new fake proteome? (y/n): " ).lower ().startswith ('y' )
397-
398- if not generate_new :
399- print ("Cannot proceed without a proteome. Exiting." )
400- sys .exit (1 )
401-
402- # Configure proteome generation (no reference needed)
403- num_proteins , target_lengths = configure_proteome_generation ()
405+ print (f"Error: Unsupported LLM model '{ args .llm_model } , file = sys .stderr )
406+ sys .exit (1 )
407+
408+ # Only run interactive proteome workflow for ProtGPT2
409+ if args .llm_model .lower () == 'protgpt2' :
410+ # Check if user has existing proteome
411+ has_existing = get_user_input ("\n Do you have an existing ProtGPT2-generated proteome file? (y/n): " ).lower ().startswith ('y' )
404412
405- # Generate the fake proteome
406- fake_proteins = generate_fake_proteome (num_proteins , target_lengths , args .llm_model )
413+ if has_existing :
414+ proteome_path = get_existing_proteome_path ()
415+ print (f"Using existing proteome: { proteome_path } )
416+ else :
417+ # Ask if user wants to generate new proteome
418+ generate_new = get_user_input ("Would you like to generate a new fake proteome? (y/n): " ).lower ().startswith ('y' )
419+
420+ if not generate_new :
421+ print ("Cannot proceed without a proteome. Exiting." )
422+ sys .exit (1 )
423+
424+ # Configure proteome generation (no reference needed)
425+ num_proteins , target_lengths = configure_proteome_generation ()
426+
427+ # Generate the fake proteome
428+ fake_proteins = generate_fake_proteome (num_proteins , target_lengths , args .llm_model )
429+
430+ # Save the generated proteome
431+ proteome_output = Path (f'fake_proteome_{ num_proteins } )
432+ write_fasta (fake_proteins , proteome_output , prefix = "protein" )
433+ print (f"\n Generated fake proteome saved to: { proteome_output } )
434+ proteome_path = proteome_output
407435
408- # Save the generated proteome
409- proteome_output = Path (f'fake_proteome_{ num_proteins } )
410- write_fasta (fake_proteins , proteome_output , prefix = "protein" )
411- print (f"\n Generated fake proteome saved to: { proteome_output } )
412- proteome_path = proteome_output
413-
414- # Now sample peptides from the proteome using FASTA method
415- print (f"\n Sampling { args .count } )
416- peptides = sample_peptides_from_fasta (proteome_path , args .length , args .count )
436+ # Now sample peptides from the proteome using FASTA method
437+ print (f"\n Sampling { args .count } )
438+ peptides = sample_peptides_from_fasta (proteome_path , args .length , args .count )
417439 else :
418440 print (f"Unknown source: { args .source } , file = sys .stderr )
419441 sys .exit (1 )
0 commit comments