gnudevx/Recommendation-System

1from bson import ObjectId2import os3from sentence_transformers import SentenceTransformer4from pymongo import MongoClient5from dotenv import load_dotenv6 7model = SentenceTransformer("all-MiniLM-L6-v2")8 9load_dotenv()10 11mongo_url = os.getenv("MONGO_URL")12 13try:14    client = MongoClient(mongo_url)15    db = client["ITJOBS"]16    jobs = db.jobs17    skills = db.skills18except Exception as e:19    print(f"Error connecting to MongoDB: {e}")20    raise RuntimeError("Failed to connect to MongoDB. Check MONGO_URL and your connection.")21 22def get_skill_names(skill_ids):23    obj_ids = [ObjectId(s) for s in skill_ids if ObjectId.is_valid(s)]24    docs = skills.find({"_id": {"$in": obj_ids}})25    return [doc["name"] for doc in docs]26 27 28def merge_skills(job):29    """Tạo skills nếu job không có trường này."""30    if "skills" in job and isinstance(job["skills"], list) and len(job["skills"]) > 0:31        return job["skills"]32 33    merged = []34 35    for field in ["mustHaveSkills", "optionalSkills", "domainKnowledge", "languages"]:36        value = job.get(field, [])37        if isinstance(value, list):38            merged.extend(value)39 40    return merged41 42 43def update_all_job_embeddings(batch_size=50):44    """45    Generate embeddings for ALL jobs, không phân biệt đã có hay chưa46    """47    # Lấy tất cả jobs48    query = {}49    50    total = jobs.count_documents(query)51    print(f"Total jobs in database: {total}")52    print(f"Generating embeddings for ALL {total} jobs...\n")53 54    # Lấy tất cả jobs, sắp xếp mới nhất trước55    all_jobs = jobs.find(query).sort("createdAt", -1)56 57    processed = 058    batch = []59 60    for job in all_jobs:61        batch.append(job)62        if len(batch) >= batch_size:63            process_batch(batch)64            processed += len(batch)65            print(f"Processed {processed}/{total} jobs...")66            batch = []67 68    if batch:69        process_batch(batch)70        processed += len(batch)71        print(f"Processed {processed}/{total} jobs...")72 73    print(f"\n✔ Generated embeddings for {processed} jobs!")74    75    # Verify final count76    final_with_embedding = jobs.count_documents({"embedding": {"$exists": True}})77    final_without = jobs.count_documents({"embedding": {"$exists": False}})78    print(f"Final: {final_with_embedding} jobs with embedding, {final_without} without")79 80 81def process_batch(job_batch):82    for job in job_batch:83        merged_skills = merge_skills(job)84 85        # Xử lý skill name86        if merged_skills and all(ObjectId.is_valid(s) for s in merged_skills):87            skill_names = get_skill_names(merged_skills)88        else:89            skill_names = merged_skills90 91        skills_text = " ".join(skill_names)92 93        req = " ".join(job.get("requirements", [])) if isinstance(job.get("requirements"), list) else job.get("requirements", "")94        desc = job.get("description", "")95 96        emb_req = model.encode(req or "")97        emb_desc = model.encode(desc or "")98        emb_skill = model.encode(skills_text or "")99 100        job_emb = 0.5 * emb_req + 0.3 * emb_desc + 0.2 * emb_skill101 102        update_data = {103            "embedding": job_emb.tolist(),104            "skills": merged_skills105        }106 107        jobs.update_one(108            {"_id": job["_id"]},109            {"$set": update_data}110        )111 112 113if __name__ == "__main__":114    print("running main...")115    update_all_job_embeddings()116