app.py

# التثبيت (لو مش مثبت)
!pip install rdkit-pypi py3Dmol transformers selfies biopython gradio -q

# الاستيراد
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 gradio as gr

# إعداد العشوائية
torch.manual_seed(int(time.time()))
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(int(time.time()))

# تحميل النموذج
model_name = "ncfrey/ChemGPT-1.2B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# تحميل ملف PDB
def load_pdb(file_obj):
    parser = PDB.PDBParser(QUIET=True)
    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()

# تنظيف وتحويل SELFIES إلى SMILES
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'])]
    cleaned_selfies = ''.join(valid_tokens)
    try:
        smiles = sf.decoder(cleaned_selfies)
        mol = Chem.MolFromSmiles(smiles)
        if mol:
            return Chem.MolToSmiles(mol)
    except:
        return None

# توليد SMILES
def generate_multiple_valid_smiles(prompt, n=10, max_length=100):
    valid_smiles = set()
    tries = 0
    while len(valid_smiles) < n and tries < n * 5:
        inputs = tokenizer(prompt, return_tensors="pt")
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            do_sample=True,
            temperature=1.0,
            top_k=100,
            pad_token_id=tokenizer.eos_token_id
        )
        selfies_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
        smiles = clean_and_decode_selfies(selfies_output)
        if smiles:
            valid_smiles.add(smiles)
        tries += 1
    return list(valid_smiles)

# الوظيفة الرئيسية
def generate_from_pdb(pdb_file):
    try:
        if isinstance(pdb_file, (str, bytes)):
            pdb_str = pdb_file if isinstance(pdb_file, str) else pdb_file.decode('utf-8', errors='ignore')
        else:
            pdb_bytes = pdb_file.read()
            pdb_str = pdb_bytes.decode('utf-8', errors='ignore')

        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)

        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 "✅ تم توليد المركبات بنجاح", html_3d, smiles_file_path

    except Exception as e:
        return f"❌ حدث خطأ:\n{str(e)}", None, None

# CSS لتجميل الواجهة
css = """
body {background-color: #f0f9ff;}
h1 {color: #004d66; text-align: center;}
"""

# واجهة Gradio
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.")
    with gr.Row():
        pdb_input = gr.File(label="📁 Upload PDB File")
        run_btn = gr.Button("🚀 Generate Compounds")
    status = gr.Textbox(label="📢 Status")
    view3d = gr.HTML(label="🧬 3D Structure Viewer")
    file_output = gr.File(label="📄 Download SMILES File")
    run_btn.click(
        fn=generate_from_pdb,
        inputs=pdb_input,
        outputs=[status, view3d, file_output]
    )

demo.launch(share=True)  # خليها True لو عايز لينك عام