# app.py

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    AutoModelForTokenClassification,
    pipeline
)
import gradio as gr
import torch


# Use a pipeline as a high-level helper
from transformers import pipeline


CLASS_MODEL_NAME = "AmandaCAI/resume-classifier"  
NER_MODEL_NAME = "AmandaCAI/ner-keywords-extract"     

# 初始化分类模型
class_tokenizer = AutoTokenizer.from_pretrained(CLASS_MODEL_NAME)
class_model = AutoModelForSequenceClassification.from_pretrained(CLASS_MODEL_NAME)

# 初始化NER模型
ner_tokenizer = AutoTokenizer.from_pretrained(NER_MODEL_NAME)
ner_model = AutoModelForTokenClassification.from_pretrained(NER_MODEL_NAME)
ner_pipeline = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple")

# 岗位分类标签（根据你的训练数据调整）
job_categories = [
    "Data Science", "Java Developer", "HR", 
    "Python Developer", "Web Designing", "Testing"
]

def analyze_resume(text):
    """处理简历分析的主函数"""
    # 岗位分类
    class_inputs = class_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        class_logits = class_model(**class_inputs).logits
    predicted_class = torch.argmax(class_logits).item()
    class_label = job_categories[predicted_class]
    
    # 技能提取
    ner_results = ner_pipeline(text)
    skills = [entity["word"] for entity in ner_results if entity["entity_group"] == "SKILL"]
    
    # 工作经验提取（示例）
    experience = [entity["word"] for entity in ner_results if entity["entity_group"] == "EXPERIENCE"]
    
    return {
        "岗位类别": class_label,
        "匹配度": f"{torch.softmax(class_logits, dim=1)[0][predicted_class].item()*100:.1f}%",
        "核心技能": list(set(skills))[:5],  # 取前5个不重复技能
        "工作经验": list(set(experience))[:3]
    }

# Gradio界面设计
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧠 AI Resume Analyzer")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="📝 Paste Resume Text Here", lines=10, 
                                  placeholder="Enter your resume text here...")
            submit_btn = gr.Button("Start the analysis", variant="primary")
            
        with gr.Column():
            output_json = gr.JSON(label="Analysis result")
            
    # 示例数据
    gr.Examples(
        examples=[[
            """John Smith
            Senior Python Developer
            Skills: Python, Django, AWS, Machine Learning
            Experience: 5+ years at Google, 3 years at Amazon
            Education: MIT Computer Science"""
        ]],
        inputs=[input_text]
    )

    submit_btn.click(
        fn=analyze_resume,
        inputs=[input_text],
        outputs=output_json
    )

# 启动应用
if __name__ == "__main__":
    demo.launch()