Update app.py

app.py

@@ -99,16 +99,17 @@ if st.button('Generate Molecules'):
     # Tokenize the core SMILES
     input_ids = tokenizer(core_smiles, return_tensors='pt').input_ids
     
-    # Generate molecules using the model
+    # Generate molecules using the model with diverse beam search
     output_ids = model.generate(
-        input_ids=input_ids,
-        max_length=128,
-        temperature=creativity,
-        do_sample=True,
-        top_k=50,
-        num_return_sequences=num_molecules,
-        num_beams = max(num_molecules, 5) # num_beams=1 was taking too long
+    input_ids=input_ids,
+    max_length=128,
+    do_sample=False,  # Use deterministic beam search
+    num_return_sequences=num_molecules,
+    num_beams=max(num_molecules * 2, 5),  # Ensure enough beams for diversity
+    num_beam_groups=num_molecules,  # Equal to num_return_sequences
+    diversity_penalty=0.5  # Adjust this value as needed
     )
+
     
     # Decode generated molecule SMILES
     generated_smiles = [