semantic-apollo/semantic_filter.py at master · uncle-rennus/semantic-apollo · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

import argparse

import pandas as pd

import yaml

from sentence_transformers import SentenceTransformer, util


def load_config(config_path="config.yaml"):

    """Loads the YAML configuration file."""

    with open(config_path, 'r', encoding='utf-8') as f:

        return yaml.safe_load(f)


def get_persona_text(config, persona_name):

    """Extracts and combines persona description and keywords."""

    persona = config['personas'].get(persona_name)

    if not persona:

        raise ValueError(f"Persona '{persona_name}' not found in config.yaml")


    description = persona.get('description', '')

    keywords = []

    for category in persona.get('keywords', {}).values():

        keywords.extend(category)


    return description + " " + " ".join(keywords)


def find_best_fit(persona_name, model_name, input_csv, config_path, output_filename=None):

    """

    Finds the best fit between a persona and profiles in a CSV file.

    """

    config = load_config(config_path)


    # 1. Load Sentence-BERT model

    print(f"Loading model: {model_name}...")

    model = SentenceTransformer(model_name)


    # 2. Get persona text and create its embedding

    persona_text = get_persona_text(config, persona_name)

    print(f"Generating embedding for persona: {persona_name}...")

    persona_embedding = model.encode(persona_text, convert_to_tensor=True)


    # 3. Load and process the enriched profiles

    print(f"Loading profiles from: {input_csv}...")

    profiles_df = pd.read_csv(input_csv)


    # Combine relevant text fields from the profile

    profiles_df['profile_text'] = profiles_df['headline'].fillna('') + ' ' + \

                                  profiles_df['title'].fillna('') + ' ' + \

                                  profiles_df['summary'].fillna('') + ' ' + \

                                  profiles_df['skills'].fillna('') + ' ' + \

                                  profiles_df['employment_description'].fillna('') + ' ' + \

                                  profiles_df['organization_keywords'].fillna('') + ' ' + \

                                  profiles_df['organization_industry'].fillna('') + ' ' + \

                                  profiles_df['past_employment_titles'].fillna('') + ' ' + \

                                  profiles_df['past_employment_descriptions'].fillna('') + ' ' + \

                                  profiles_df['contact_headline'].fillna('')


    # 4. Generate embeddings for all profiles

    print("Generating embeddings for all profiles...")

    profile_embeddings = model.encode(profiles_df['profile_text'].tolist(), convert_to_tensor=True)


    # 5. Calculate cosine similarity

    print("Calculating similarity scores...")

    cosine_scores = util.pytorch_cos_sim(persona_embedding, profile_embeddings)


    # 6. Add scores to the DataFrame and save

    profiles_df['similarity_score'] = cosine_scores[0].tolist()


    # Sort by score and select top results

    results_df = profiles_df.sort_values(by='similarity_score', ascending=False)


    if output_filename is None:

        output_filename = f"output/{persona_name}_semantic_fit.csv"


    results_df.to_csv(output_filename, index=False)


    print(f"Results saved to {output_filename}")

    print("\nTop 5 matches:")

    print(results_df[['linkedin_url', 'headline', 'title', 'similarity_score']].head(5))


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Semantic profile filtering with Sentence-BERT.")

    parser.add_argument("--persona", required=True, help="The name of the persona to use for matching.")

    parser.add_argument("--model", default="paraphrase-multilingual-MiniLM-L12-v2", help="The Sentence-BERT model to use.")

    parser.add_argument("--input", default="apollo_people_data_enriched.csv", help="The input CSV file with enriched profiles.")

    parser.add_argument("--config", default="config.yaml", help="The configuration YAML file.")

    parser.add_argument("--output", help="The output CSV file path.")


    args = parser.parse_args()


    find_best_fit(args.persona, args.model, args.input, args.config, args.output)