import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os
import time

# ====================== 核心配置（CPU极致优化）======================
MODEL_NAME = "PharMolix/BioMedGPT-LM-7B"
MAX_NEW_TOKENS = 200  # 进一步降低，避免CPU过载（免费版推荐150-200）
TEMPERATURE = 0.6
MAX_INPUT_LENGTH = 350  # 缩短输入长度，节省内存
DEVICE = "cpu"

# 加速下载+避免缓存冲突
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["TRANSFORMERS_CACHE"] = "./cache"

# ====================== 模型加载（稳定性优化）======================
print(f"正在加载模型：{MODEL_NAME}（首次部署约10-15分钟）...")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    padding_side="right",
    truncation_side="left",
    legacy=False
)

# 强制设置pad_token（双重保险）
if not tokenizer.pad_token:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id
print(f"Pad_token设置完成：{tokenizer.pad_token}（ID：{tokenizer.pad_token_id}）")

# 模型加载（CPU推理极致优化）
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map=DEVICE,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
    torch_dtype=torch.float32,
    offload_folder="./offload",
    offload_state_dict=True,
    max_memory={DEVICE: "10GB"}  # 限制CPU内存使用（适配免费环境）
)

# 推理模式（比torch.no_grad()更省资源）
model.eval()
print("模型加载完成！免费CPU生成单条回复约20-60秒，请耐心等待～")

# ====================== 核心问答逻辑（状态反馈+错误捕获）======================
def biomed_qa(prompt, task_type="医学问答", progress=gr.Progress(track_tqdm=True)):
    """
    带进度反馈+异常捕获的问答函数
    progress：Gradio进度条，实时反馈状态
    """
    try:
        # 1. 进度反馈：开始处理
        progress(0.1, desc="正在解析问题...")
        time.sleep(1)  # 避免进度条闪跳

        # 2. 构建专业Prompt（确保模型理解任务）
        task_prompt = {
            "医学问答": "请从生物医学专业角度详细解答，逻辑清晰，分点说明核心机制、临床应用、潜在风险：\n",
            "药物机制": "请分析以下药物的作用靶点、分子机制、代谢途径及临床适应症，基于公开研究依据简要说明：\n",
            "文献解读": "请简洁解读以下生物医学文献的核心结论、关键实验方法及临床转化价值：\n",
            "靶点预测": "基于药物名称/SMILES序列，预测1-2个核心作用靶点及简要结合模式，说明预测依据：\n"
        }
        full_prompt = task_prompt[task_type] + prompt
        progress(0.2, desc="正在编码输入...")

        # 3. 输入编码（CPU优化：关闭padding，限制长度）
        inputs = tokenizer(
            full_prompt,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_INPUT_LENGTH,
            padding=False
        )
        progress(0.3, desc="模型推理中（免费CPU约20-60秒）...")

        # 4. 模型生成（关键优化：禁用梯度+推理模式，避免内存溢出）
        with torch.inference_mode():  # 比torch.no_grad()更省CPU资源
            outputs = model.generate(
                **inputs,
                max_new_tokens=MAX_NEW_TOKENS,
                temperature=TEMPERATURE,
                top_p=0.8,
                repetition_penalty=1.3,
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                num_return_sequences=1,
                no_repeat_ngram_size=3,
                early_stopping=True  # 遇到eos_token立即停止，避免无效生成
            )
        progress(0.8, desc="正在整理回复...")

        # 5. 解码+清理回复（确保内容有效）
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # 移除输入Prompt前缀，只保留模型生成部分
        clean_response = response.replace(full_prompt, "").strip()

        # 6. 处理空回复情况
        if not clean_response or len(clean_response) < 10:
            clean_response = "⚠️  模型生成失败（可能是输入过长或CPU资源不足），请尝试：\n1. 简化问题或缩短输入长度\n2. 等待当前生成完成后再提问\n3. 选择更简洁的任务类型（如靶点预测）"

        # 7. 最终回复（带提示）
        final_response = f"{clean_response}\n\n⏱️  生成耗时：{time.time() - start_time:.1f}秒\n⚠️  提示：本回答仅作科研参考，不构成医疗/用药建议"
        progress(1.0, desc="生成完成！")
        return final_response

    except Exception as e:
        # 捕获所有异常，显示到界面（关键！避免静默失败）
        error_msg = f"❌ 生成出错：{str(e)[:500]}（完整错误见Space日志）\n\n建议解决方案：\n1. 关闭当前页面重新打开\n2. 简化问题，减少输入字数\n3. 等待30秒后再尝试\n4. 更换任务类型（如靶点预测比文献解读更省资源）"
        return error_msg

# ====================== Gradio界面（强化交互反馈）======================
with gr.Blocks(title="BioMedGPT 生物医药助手（免费CPU版）") as demo:
    gr.Markdown("""
    # 🧬 BioMedGPT 生物医药智能助手（免费CPU版）
    ## 🔔 重要提示
    1. 免费CPU性能有限，单条回复生成约 **20-60秒**，请耐心等待（点击按钮后勿重复点击）
    2. 建议输入长度≤100字，优先选择「靶点预测」「药物机制」等简洁任务
    3. 若生成失败，可简化问题或等待1分钟后再试
    ## 🔬 核心功能
    - 医学问答：基础/临床问题专业解答
    - 药物机制：靶点、机制、临床应用分析
    - 文献解读：核心结论+转化价值提炼
    - 靶点预测：药物-SMILES靶点预测
    """)
    
    # 输入区域
    with gr.Row(equal_height=True):
        task_type = gr.Dropdown(
            choices=["医学问答", "药物机制", "文献解读", "靶点预测"],
            value="药物机制",  # 默认选更省资源的任务
            label="任务类型",
            info="优先选择「药物机制」「靶点预测」，生成更快"
        )
    prompt = gr.Textbox(
        label="输入问题/内容",
        placeholder="示例1：PD-1抑制剂的作用靶点是什么？\n示例2：阿司匹林（SMILES: CC(=O)OC1=CC=CC=C1C(=O)O）的核心靶点？",
        lines=3,
        max_lines=4,
        max_length=150  # 限制输入长度，避免CPU过载
    )
    
    # 输出区域（增加边框+复制功能）
    output = gr.Textbox(
        label="专业解答",
        lines=12,
        max_lines=20,
        show_copy_button=True,
        container=True,
        elem_id="response-box",
        interactive=False
    )
    
    # 交互按钮（禁用重复点击）
    btn = gr.Button("生成专业解答", variant="primary", size="lg")
    # 绑定按钮事件，添加进度条反馈
    btn.click(
        fn=biomed_qa,
        inputs=[prompt, task_type],
        outputs=output,
        show_progress="full"  # 显示完整进度条
    )
    
    # 快速示例（选择更简洁的问题，确保能生成）
    with gr.Row(equal_height=True, elem_id="examples"):
        gr.Button("示例1：PD-1靶点").click(
            lambda: ("PD-1抑制剂的核心作用靶点及机制？", "药物机制"),
            outputs=[prompt, task_type]
        )
        gr.Button("示例2：阿司匹林靶点").click(
            lambda: ("阿司匹林（SMILES: CC(=O)OC1=CC=CC=C1C(=O)O）的核心靶点？", "靶点预测"),
            outputs=[prompt, task_type]
        )
    
    # 清空功能
    gr.Button("清空输入输出", variant="secondary").click(
        lambda: ("", "药物机制", ""),
        outputs=[prompt, task_type, output]
    )

# ====================== 启动服务（适配免费环境）======================
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,  # 界面显示错误
        debug=False,
        inline=False
    )