Update app.py

app.py

@@ -1,12 +1,11 @@
 from Bio import PDB
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from rdkit import Chem
-import py3Dmol
-import re
-import io
 import selfies as sf
 import torch
 import time
+import re
+import io
 import gradio as gr
 
 # إعداد العشوائية
@@ -14,7 +13,7 @@ torch.manual_seed(int(time.time()))
 if torch.cuda.is_available():
     torch.cuda.manual_seed_all(int(time.time()))
 
-# تحميل النموذج
+# تحميل نموذج ChemGPT
 model_name = "ncfrey/ChemGPT-1.2B"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
@@ -24,13 +23,6 @@ def load_pdb(file_obj):
     structure = parser.get_structure('protein', file_obj)
     return structure
 
-def get_protein_3d_html(pdb_str):
-    view = py3Dmol.view(width=600, height=400)
-    view.addModel(pdb_str, "pdb")
-    view.setStyle({"cartoon": {"color": "spectrum"}})
-    view.zoomTo()
-    return view._make_html()
-
 def clean_and_decode_selfies(raw_output):
     tokens = re.findall(r'\[[^\[\]]+\]', raw_output)
     valid_tokens = [t for t in tokens if all(x not in t for x in ['Branch', 'Ring', 'expl'])]
@@ -63,58 +55,33 @@ def generate_multiple_valid_smiles(prompt, n=10, max_length=100):
         tries += 1
     return list(valid_smiles)
 
-def generate_from_pdb(pdb_file):
+def generate_drugs_from_pdb(pdb_file):
     try:
-        with open(pdb_file.name, 'r', encoding='utf-8', errors='ignore') as f:
-            pdb_str = f.read()
-
-        if len(pdb_str.strip()) == 0:
-            return "❌ الملف فارغ أو غير صالح", None, None
-
-        pdb_file_io = io.StringIO(pdb_str)
-        try:
-            load_pdb(pdb_file_io)
-        except Exception as e:
-            return f"❌ خطأ أثناء تحليل ملف PDB:\n{str(e)}", None, None
-
-        html_3d = get_protein_3d_html(pdb_str)
+        pdb_bytes = pdb_file.read()
+        pdb_str = pdb_bytes.decode('utf-8', errors='ignore')
+        pdb_io = io.StringIO(pdb_str)
+        load_pdb(pdb_io)
 
         prompt = "Generate a molecule in SELFIES that binds to the mutated KRAS protein"
         smiles_list = generate_multiple_valid_smiles(prompt, n=10)
 
         if not smiles_list:
-            return "❌ لم يتم توليد أي SMILES صالحة", html_3d, None
-
-        smiles_txt = "\n".join(smiles_list)
-        smiles_file_path = "/tmp/generated_smiles.txt"
-        with open(smiles_file_path, "w") as f:
-            f.write(smiles_txt)
+            return "❌ لم يتم توليد أي SMILES صالحة", ""
 
-        return "✅ تم توليد المركبات بنجاح", html_3d, smiles_file_path
+        smiles_text = "\n".join(smiles_list)
+        return "✅ تم توليد المركبات بنجاح", smiles_text
 
     except Exception as e:
-        return f"❌ حدث خطأ:\n{str(e)}", None, None
-
-css = """
-body {background-color: #f0f9ff;}
-h1 {color: #004d66; text-align: center;}
-"""
-
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("<h1>🔬 Drug-like Molecule Generation from PDB using ChemGPT</h1>")
-    gr.Markdown("🧪 Upload a PDB file containing mutations in the KRAS protein. The system will generate suitable SMILES drug candidates.")
+        return f"❌ خطأ: {str(e)}", ""
 
+# واجهة Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧬 توليد مركبات دوائية من ملف PDB باستخدام ChemGPT")
     with gr.Row():
-        pdb_input = gr.File(label="📁 Upload PDB File")
-        run_btn = gr.Button("🚀 Generate Molecules")
-    status = gr.Textbox(label="📢 Status")
-    view3d = gr.HTML(label="🧬 3D Structure View")
-    file_output = gr.File(label="📄 Download SMILES File")
-    run_btn.click(
-        fn=generate_from_pdb,
-        inputs=pdb_input,
-        outputs=[status, view3d, file_output]
-    )
+        pdb_input = gr.File(label="📁 ارفع ملف PDB")
+        run_btn = gr.Button("🚀 توليد SMILES")
+    status = gr.Textbox(label="📢 الحالة")
+    smiles_output = gr.Textbox(label="📄 المركبات (SMILES)", lines=10)
+    run_btn.click(fn=generate_drugs_from_pdb, inputs=pdb_input, outputs=[status, smiles_output])
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()