import gradio as gr
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    AutoModelForSeq2SeqLM,
)
import torch
model_folder = "zjunlp/chatcell-large"
tokenizer = AutoTokenizer.from_pretrained(model_folder)
model = AutoModelForSeq2SeqLM.from_pretrained(model_folder)
if torch.cuda.is_available():
    device = torch.device("cuda")
    model.to(device)  
else:
    device = torch.device("cpu")

model.eval()


def run_detector(input_text):
    # Encode the input text and generate a response with specified generation parameters
    input_ids = tokenizer(input_text,return_tensors="pt").input_ids.to(device)
    output_ids = model.generate(input_ids, max_length=512, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50, top_p=0.95, do_sample=True)
    # Decode and print the generated output text
    output_text = tokenizer.decode(output_ids[0],skip_special_tokens=True)

    return output_text




css = """
.green { color: black!important;line-height:1.9em; padding: 0.2em 0.2em; background: #ccffcc; border-radius:0.5rem;}
.red { color: black!important;line-height:1.9em; padding: 0.2em 0.2em; background: #ffad99; border-radius:0.5rem;}
.hyperlinks {
  display: flex;
  align-items: center;
  align-content: center;
  padding-top: 12px;
  justify-content: flex-end;
  margin: 0 10px; /* Adjust the margin as needed */
  text-decoration: none;
  color: #000; /* Set the desired text color */
}
"""

capybara_problem=''
with gr.Blocks(css=css,
               theme=gr.themes.Soft(text_size="sm")) as app:
    
    with gr.Row():
        gr.HTML("""
                <div style="text-align: center; margin: 0 auto;">
                <p><h1> ChatCell: Facilitating Single-Cell Analysis with Natural Language</h1>
                </div>

        """)
        # gr.HTML("")
    with gr.Row():
        gr.Markdown("<p align='center'><a href='https://github.com/zjunlp/ChatCell'>🐣https://github.com/zjunlp/ChatCell</a></p>")
    with gr.Row():
        input_box = gr.Textbox(value=capybara_problem, placeholder="Enter text here", lines=4, label="Input Text", )
    with gr.Row():
        output_text = gr.Textbox(label="Prediction")

    with gr.Row():
        clear_button = gr.ClearButton()
        submit_button = gr.Button(variant="primary")
    # with gr.Column(scale=2):

    examples = gr.Examples(
        examples=[
            ["Distinguish between resistant and sensitive cancer cells in response to Cisplatin, using the data from the 100 most expressed genes in descending order MYL12B FTL MYL12A HIST1H4C RPL23 GSTP1 RPS3 ENO1 RPLP1 TXN ANXA2 PPP1CB B2M RPLP0 HSPA8 H2AFZ TPI1 ANXA1 RPL7 GAPDH CHP1 LDHA RPL3 S100A11 PRDX1 CALM2 CAPZA1 SLC25A5 RPS27 YWHAZ GNB2L1 PTBP3 RPS6 MOB1A S100A2 ACTG1 BROX SAT1 RPL35A CA2 PSMB4 RPL8 TBL1XR1 RPS18 HNRNPH1 RPL27 RPS14 RPS11 ANP32E RPL19 C6ORF62 RPL9 EEF1A1 RPL5 COLGALT1 NPM1 CCT6A RQCD1 CACUL1 RPL4 HSP90AA1 MALAT1 ALDOA PSMA4 SEC61G RPL38 PSMB5 FABP5 HSP90AB1 RPL35 CHCHD2 EIF3E COX4I1 RPL21 PAFAH1B2 PTMA TMED4 PSMB3 H3F3B AGO1 DYNLL1 ATP5A1 LDHB COX7B ACTB RPS27A PSME2 ELMSAN1 NDUFA1 HMGB2 PSMB6 TMSB10 SET RPL12 RPL37A RPS13 EIF1 ATP5G1 RPS3A TOB1."],
            ["Evaluate a cancer cell's response to Erlotinib (resistant or sensitive), based on the cell's 100 most actively expressed genes in descending order MT-RNR2 MT-CO3 MT-CO1 RPL13 RPLP1 FTH1 GAPDH RPS8 PABPC1 FTL ANXA2 NCL RPS12 PTMA RPS14 CAV1 RPS3 RPS4X MT-CO2 RPL37 KRT7 RPS2 TMSB4X HSP90AA1 GSTP1 MT-ND4 MT-ATP6 S100A6 NAP1L1 RPL31 HSP90AB1 B2M RPL19 PPIA NUCKS1 MT-ND5 RPS27A TPM1 RPL18 ATP5G3 RPLP0 RPL8 TXN GNAS PSMA7 RPL30 MYL6 SLC25A5 RAD21 RTN4 RPL37A HSPD1 LRRFIP1 DEK MT-ND1 RPL11 TPT1 TMSB10 RPS24 SSRP1 LUC7L3 RPS20 KRT18 RPS16 RPL5 RPL35 RPS21 HNRNPC RPLP2 NACA GNB2L1 LGALS1 UQCRH FAU CHCHD2 RPL23 LDHB UBA52 HSP90B1 SEC62 RPL6 DSTN RPL27A RPL23A RPL22L1 KTN1 RPL14 CALM2 PRDX1 ADIPOR2 ZFAS1 HIST1H4C UQCRQ CALU PTTG1 EIF1 RPL26 RPL10A RAN ARPC2."]
        ],
        examples_per_page=3,
        inputs=[input_box],
    )



    with gr.Accordion("Disclaimer", open=False):
        gr.Markdown(
            """
            - `Accuracy` :
                - This model aims for accuracy but cannot guarantee 100% precision. Results should be used as a guide, not as definitive tools.
            - `Use Restrictions`: 
                - Intended for educational and research purposes only. Not for clinical or commercial use without prior verification.
            - `Responsibility`: 
                - Users are solely responsible for any outcomes resulting from the use of this model.             
            """
        )
    with gr.Accordion("Cite our work", open=False):
        gr.Markdown(
            """
            ```bibtex
            @article{fang2024chatcell,
                  title={ChatCell: Facilitating Single-Cell Analysis with Natural Language},
                  author={Fang, Yin and Liu, Kangwei and Zhang, Ningyu and Deng, Xinle and Yang, Penghui and Chen, Zhuo and Tang, Xiangru and Gerstein, Mark and Fan, Xiaohui and Chen, Huajun},
                  year={2024}
            }

            """
        )
    submit_button.click(run_detector, inputs=input_box, outputs=output_text)
    clear_button.click(lambda: ("", ""), outputs=[input_box, output_text])

app.launch()
# 📙