app.py

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 设置模型路径（替换为你上传后显示的路径名）
model_path = "tosei0000/code"  # 修改为你上传的文件夹名

# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)

# 如果你有 GPU（Kaggle 支持 GPU），把模型放到 GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# 推理函数
def chat(prompt, max_new_tokens=100):
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 测试
response = chat("你好，请介绍一下你自己。")
print(response)