ai_mapper.py

# ai_mapper.py

from sentence_transformers import SentenceTransformer, util

# Load model once at module level
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

def ai_map_fields(extracted_keys, object_fields):
    try:
        mappings = {}
        confidence_scores = {}

        field_embeddings = model.encode(object_fields, convert_to_tensor=True)

        for key in extracted_keys:
            if key.lower() in ["name", "email"]:
                continue

            key_embedding = model.encode(key, convert_to_tensor=True)
            cosine_scores = util.pytorch_cos_sim(key_embedding, field_embeddings)[0]
            best_score_idx = cosine_scores.argmax().item()
            best_match = object_fields[best_score_idx]
            confidence = cosine_scores[best_score_idx].item()

            mappings[key] = best_match
            confidence_scores[key] = round(confidence, 2)

        return mappings, confidence_scores, None
    except Exception as e:
        return None, None, str(e)