Update app.py

app.py

@@ -17,15 +17,22 @@ def drug_discovery(disease, 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}. Output only the SMILES."
+    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@+\-\[\]\(\)=#$]{5,}", smiles_result)
-    smiles = smiles_match[0] if smiles_match else "CC(C)CC"  # fallback
+    smiles_match = 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:
+        mol_test = Chem.MolFromSmiles(match)
+        if mol_test:
+            smiles = match
+            break
+    if not smiles:
+        smiles = "CC(C)CC"  # fallback to isopentane
 
     # Generate RDKit molecule
     mol = Chem.MolFromSmiles(smiles)
     if not mol:
-        return "Invalid SMILES generated", "", "", ""
+        return "Invalid SMILES generated", smiles, "", ""
 
     AllChem.Compute2DCoords(mol)
 
@@ -54,7 +61,6 @@ def drug_discovery(disease, symptoms):
     viewer.setStyle({"stick": {"colorscheme": "cyanCarbon"}})
     viewer.setBackgroundColor("black")
     viewer.zoomTo()
-    viewer.rotate(1)
     viewer.spin(True)
     viewer_html_raw = viewer._make_html()
 
@@ -118,3 +124,4 @@ iface = gr.Interface(
 )
 
 iface.launch(share=True)
+