import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# 加载模型
MODEL_REPO = "jinv2/ai-job-navigator-model"
base_model = AutoModelForCausalLM.from_pretrained("distilgpt2")
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
model = PeftModel.from_pretrained(base_model, MODEL_REPO)

# 定义生成函数
def generate_advice(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_length=200,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.9,
        num_return_sequences=1
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# 创建 Gradio 界面
interface = gr.Interface(
    fn=generate_advice,
    inputs=gr.Textbox(label="输入提示", placeholder="根据最新的AI行业趋势，提供2025年的职业建议："),
    outputs=gr.Textbox(label="生成结果"),
    title="AI Job Navigator 2025",
    description="输入提示以获取 2025 年 AI 行业的职业建议（基于微调的 distilgpt2 模型）。注意：由于训练数据有限，生成结果可能不理想。"
)

interface.launch()