import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Đường dẫn tới mô hình trên Hugging Face Hub
model_name = "trangannh/ptit-job-recommendation"

# Khởi tạo tokenizer và mô hình từ tên mô hình trên Hugging Face
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Hàm dự đoán công việc dựa trên kỹ năng cứng, kỹ năng mềm và ngành nghề
def recommend_jobs(input_hard_skills, input_soft_skills, input_major, top_n=3):
    # Chuẩn bị input cho mô hình
    inputs = {
        "hard_skills": input_hard_skills,
        "soft_skills": input_soft_skills,
        "major": input_major
    }
    
    # Tiền xử lý và mã hóa input
    encoded_input = tokenizer.encode_plus(
        inputs["hard_skills"],  # Kỹ năng cứng
        inputs["soft_skills"],  # Kỹ năng mềm
        inputs["major"],        # Ngành nghề
        add_special_tokens=True,
        return_tensors="pt"
    )
    
    # Dự đoán
    with torch.no_grad():
        outputs = model(**encoded_input)
    
    # Lấy giá trị dự đoán và sắp xếp theo thứ tự giảm dần
    logits = outputs.logits[0].tolist()
    sorted_indices = sorted(range(len(logits)), key=lambda k: logits[k], reverse=True)
    
    # Lấy top N công việc gợi ý
    recommended_jobs = []
    for i in range(min(top_n, len(sorted_indices))):
        job_index = sorted_indices[i]
        recommended_jobs.append(tokenizer.decode(job_index))
    
    return recommended_jobs

# Hướng dẫn sử dụng
if __name__ == "__main__":
    # Input từ người dùng (có thể làm thay đổi để phù hợp với nhu cầu thực tế)
    input_hard_skills = input("Nhập kỹ năng cứng của bạn: ")
    input_soft_skills = input("Nhập kỹ năng mềm của bạn: ")
    input_major = input("Nhập ngành nghề của bạn: ")
    
    # Gợi ý công việc
    recommended_jobs = recommend_jobs(input_hard_skills, input_soft_skills, input_major)
    
    # In kết quả
    print(f"Các công việc được gợi ý cho bạn:")
    for i, job in enumerate(recommended_jobs, start=1):
        print(f"{i}. {job}")