app.py

import os
import gc
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer

DEFAULT_MODEL_SMALL = "vandijklab/C2S-Scale-Gemma-2-2B"
DEFAULT_MODEL_LARGE = "vandijklab/C2S-Scale-Gemma-2-27B"

MODEL_CACHE = {"id": None, "tokenizer": None, "model": None}

def vram_gb():
    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(0)
        return props.total_memory / (1024**3)
    return 0.0

def build_prompt(gene_list, species="Homo sapiens"):
    if isinstance(gene_list, str):
        # permitir lista separada por comas/espacios/nuevas líneas
        raw = [g.strip() for g in gene_list.replace("\n", ",").split(",") if g.strip()]
        genes = ", ".join(raw)
    else:
        genes = ", ".join(gene_list)
    return (
        f"The following is a list of gene names ordered by descending expression level "
        f"in a {species} cell. Your task is to give the cell type which this cell belongs "
        f"to based on its gene expression.\n"
        f"Cell sentence: {genes}.\n"
        f"The cell type corresponding to these genes is:"
    )

def unload():
    MODEL_CACHE["id"] = None
    MODEL_CACHE["tokenizer"] = None
    MODEL_CACHE["model"] = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_model(model_id, quantization):
    """
    Carga perezosa del modelo. Para 27B se recomienda A100 80GB.
    quantization: 'none' o '8bit' (requiere bitsandbytes).
    """
    if MODEL_CACHE["id"] == model_id and MODEL_CACHE["model"] is not None:
        return MODEL_CACHE["tokenizer"], MODEL_CACHE["model"]

    unload()

    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
    device_map = "auto" if torch.cuda.is_available() else {"": "cpu"}

    kwargs = dict(torch_dtype=dtype, device_map=device_map, low_cpu_mem_usage=True)

    if quantization == "8bit":
        try:
            import bitsandbytes as bnb  # noqa: F401
            kwargs.update(dict(load_in_8bit=True))
        except Exception:
            # Si no está disponible, caemos a sin cuantizar
            pass

    tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    mdl = AutoModelForCausalLM.from_pretrained(model_id, **kwargs).eval()

    MODEL_CACHE["id"] = model_id
    MODEL_CACHE["tokenizer"] = tok
    MODEL_CACHE["model"] = mdl
    return tok, mdl

def infer(model_id, species, genes_text, max_new_tokens, temperature, top_p, top_k, repetition_penalty, quantization):
    # chequeo sencillo de VRAM con guía para 27B
    mem = vram_gb()
    warn = ""
    if "27B" in model_id:
        if mem < 60 and quantization != "8bit":
            warn = (
                f"⚠️ Detectada VRAM ~{mem:.1f}GB. Para 27B se recomienda A100 80GB "
                f"o usar 8-bit (aun así puede ser insuficiente en T4)."
            )

    tok, mdl = load_model(model_id, quantization)
    prompt = build_prompt(genes_text, species=species)
    inputs = tok(prompt, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = {k: v.to(mdl.device) for k, v in inputs.items()}

    streamer = TextIteratorStreamer(tok, skip_special_tokens=True)
    gen_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=True,
        temperature=float(temperature),
        top_p=float(top_p),
        top_k=int(top_k),
        repetition_penalty=float(repetition_penalty),
        eos_token_id=tok.eos_token_id,
        streamer=streamer,
    )

    # streaming
    import threading
    output_text = ""
    def _gen():
        with torch.no_grad():
            mdl.generate(**gen_kwargs)

    thread = threading.Thread(target=_gen)
    thread.start()
    for new_text in streamer:
        output_text += new_text
        yield (warn, prompt, output_text)
    thread.join()

with gr.Blocks(title="C2S-Scale (Gemma-2) — Single-cell Biology") as demo:
    gr.Markdown(
        """
        # C2S-Scale (Gemma-2) for single-cell biology
        Infiere **tipo celular** a partir de una *cell sentence* (genes ordenados por expresión).
        - Modelos: `vandijklab/C2S-Scale-Gemma-2-2B` (ligero), `vandijklab/C2S-Scale-Gemma-2-27B` (pesado).  
        - Selecciona GPU en Settings del Space para mejor rendimiento.

        **Nota:** 27B requiere GPU grande (idealmente A100 80GB). En T4, incluso con 8-bit, puede no cargar.
        """
    )
    with gr.Row():
        model_id = gr.Dropdown(
            choices=[DEFAULT_MODEL_SMALL, DEFAULT_MODEL_LARGE],
            value=DEFAULT_MODEL_SMALL,
            label="Modelo"
        )
        quantization = gr.Radio(["none", "8bit"], value="none", label="Cuantización (experimental)")
        species = gr.Dropdown(["Homo sapiens", "Mus musculus", "Danio rerio", "Custom…"], value="Homo sapiens", label="Especie")
        species_custom = gr.Textbox(value="", label="Especie (si elegiste Custom…)", visible=False)

    def _toggle_species(choice):
        return gr.update(visible=(choice == "Custom…"))
    species.change(_toggle_species, species, species_custom)

    example_genes = "MALAT1, RPLP0, RPL13A, ACTB, RPS27A, PTPRC, CD3D, CD3E, CCR7, IL7R, LTB, TRAC, CD27, CD4, CCR6, CXCR5"
    genes_text = gr.Textbox(value=example_genes, lines=6, label="Cell sentence (lista de genes ordenados por expresión ↓)")

    with gr.Accordion("Parámetros de generación", open=False):
        max_new_tokens = gr.Slider(8, 256, value=64, step=1, label="max_new_tokens")
        temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.01, label="top_p")
        top_k = gr.Slider(1, 200, value=50, step=1, label="top_k")
        repetition_penalty = gr.Slider(0.8, 1.5, value=1.05, step=0.01, label="repetition_penalty")

    warn_box = gr.Markdown("")
    prompt_box = gr.Code(label="Prompt efectivo", language="text")
    output_box = gr.Textbox(label="Salida del modelo (stream)")

    def _species_value(sp, custom):
        return custom if sp == "Custom…" and custom.strip() else sp

    run_btn = gr.Button("🚀 Inferir tipo celular")
    run_btn.click(
        fn=lambda mid, sp, spc, genes, mx, temp, tp, tk, rp, q: infer(
            mid, _species_value(sp, spc), genes, mx, temp, tp, tk, rp, q
        ),
        inputs=[model_id, species, species_custom, genes_text, max_new_tokens, temperature, top_p, top_k, repetition_penalty, quantization],
        outputs=[warn_box, prompt_box, output_box]
    )

if __name__ == "__main__":
    demo.launch()