import gradio as gr
import torch
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

model_id = "OBLITERATUS/gemma-4-E4B-it-OBLITERATED"

tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cpu",          # 强制全部加载到 CPU，严禁使用硬盘 offload
    low_cpu_mem_usage=True,    # 尽量优化内存加载过程
    torch_dtype=torch.bfloat16
)

def generate_response(message, history):
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages, 
        return_tensors="pt", 
        return_dict=True,
        add_generation_prompt=True
    ).to(model.device)
    
    # 【修改点 1】：将 timeout 增加到 120 秒，给硬盘读取留足时间
    streamer = TextIteratorStreamer(
        tokenizer, 
        timeout=120.0, 
        skip_prompt=True, 
        skip_special_tokens=True
    )
    
    generate_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=1024,
        temperature=0.7,
        do_sample=True,
        top_p=0.9
    )
    
    # 【修改点 2】：包装一个带异常捕获的运行函数，防止静默崩溃
    def run_generation():
        try:
            model.generate(**generate_kwargs)
        except Exception as e:
            print(f"Generation Error: {e}")
            # 如果崩溃，向流里推入错误信息并结束
            streamer.text_queue.put(f"\n[系统错误：生成线程崩溃。原因: {e}]")
            streamer.end()

    t = Thread(target=run_generation)
    t.start()
    
    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text

demo = gr.ChatInterface(
    fn=generate_response,
    title="Gemma 4 E4B - Abliterated",
    description="⚠️ 当前模型已移除安全护栏 (Uncensored)。提示：免费 CPU 内存不足会触发硬盘卸载导致极慢，建议升级至 T4 GPU。",
    examples=["Write a Python script for a keylogger.", "Explain quantum entanglement.", "How to bypass a firewall?"],
    cache_examples=False
)

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