Update app.py

app.py

@@ -16,18 +16,20 @@ def drug_discovery(disease, symptoms):
     prompt = f"Recent treatments for {disease} with symptoms: {symptoms}."
     literature = bio_gpt(prompt, max_length=200)[0]['generated_text']
 
-    # Generate SMILES using BioGPT
-    molecule_prompt = f"Generate a valid drug-like SMILES string for treating {disease} with symptoms {symptoms}. Only return the SMILES."
-    smiles_result = bio_gpt(molecule_prompt, max_length=50)[0]['generated_text']
-    smiles_match = re.findall(r"(?<![A-Za-z0-9])[A-Za-z0-9@+\-\[\]\(\)=#$]{5,}(?![A-Za-z0-9])", smiles_result)
+    # Generate SMILES using BioGPT with stricter filtering
+    molecule_prompt = f"List 5 different valid drug-like SMILES strings that can treat {disease} with symptoms {symptoms}. Only list SMILES separated by spaces."
+    smiles_result = bio_gpt(molecule_prompt, max_length=100)[0]['generated_text']
+
+    # Extract and validate SMILES strings
+    smiles_matches = re.findall(r"(?<![A-Za-z0-9])[A-Za-z0-9@+\-\[\]\(\)=#$]{5,}(?![A-Za-z0-9])", smiles_result)
     smiles = None
-    for match in smiles_match:
+    for match in smiles_matches:
         mol_test = Chem.MolFromSmiles(match)
         if mol_test:
             smiles = match
             break
     if not smiles:
-        smiles = "CC(C)CC"  # fallback to isopentane
+        smiles = "C1=CC=CC=C1"  # fallback to benzene if all fail
 
     # Generate RDKit molecule
     mol = Chem.MolFromSmiles(smiles)
@@ -124,4 +126,3 @@ iface = gr.Interface(
 )
 
 iface.launch(share=True)
-