app.py

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

# --- 直接加载模型和分词器 ---
model_id = "xingyu1996/tiger-gpt2"
tokenizer = AutoTokenizer.from_pretrained("gpt2") # 使用原始的 GPT-2 分词器
model = AutoModelForCausalLM.from_pretrained(model_id)

def respond(
    message,
    history: list[tuple[str, str]],
    max_tokens,
    temperature,
    top_p,
):
    # 将输入文本转换为 token ID
    input_ids = tokenizer.encode(message, return_tensors="pt")
    
    # 准备生成参数
    gen_kwargs = {
        "max_length": input_ids.shape[1] + max_tokens,
        "do_sample": True if temperature > 0 else False,
    }
    
    if temperature > 0:
        gen_kwargs["temperature"] = temperature
    if top_p < 1.0:
        gen_kwargs["top_p"] = top_p
        
    # 生成文本
    with torch.no_grad():
        output_ids = model.generate(input_ids, **gen_kwargs)
        
    # 只保留新生成的部分
    new_tokens = output_ids[0, input_ids.shape[1]:]
    
    # 解码生成的 token ID
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
    return response


# 其他 Gradio 界面代码不变
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Slider(minimum=1, maximum=512, value=325, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
    title=f"推理测试: {model_id}",
    description="输入中文文本，模型将进行补全。"
)

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