-
Notifications
You must be signed in to change notification settings - Fork 0
/
embeddings.py
181 lines (150 loc) · 6.49 KB
/
embeddings.py
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
177
178
179
180
181
"""
Link each stereotype to a list of tags and show their similarity via word embeddings
Steps:
1. text preprocessing;
2. embedding generation for stereotypes and tags using pre-trained embeddings,
3. similarity calculation using cosine similarity,
4. visualization?
"""
import json
import logging
from typing import Dict, List
import matplotlib.pyplot as plt
import networkx as nx
import nltk
import pandas as pd
import seaborn as sns
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
# Constants for file paths
STEREOTYPES_JSON_PATH = "stereotypes.json"
TAGS_JSON_PATH = "tags.json"
PLOTS_PATH = "plots/"
# Configure logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
def download_nltk_data():
"""Download necessary NLTK data if not already present."""
nltk_data = ["stopwords", "punkt"]
for data in nltk_data:
try:
nltk.data.find(f"tokenizers/{data}")
except LookupError:
nltk.download(data)
def preprocess(text: str) -> str:
"""
Tokenize and clean the input text by removing stopwords and non-alphabetic tokens.
Parameters:
- text (str): The input text to be preprocessed.
Returns:
- str: The preprocessed text, with stopwords and non-alphabetic tokens removed.
"""
stop_words = set(stopwords.words("english"))
tokens = word_tokenize(text.lower())
return " ".join(
[word for word in tokens if word.isalpha() and word not in stop_words]
)
def load_and_preprocess_data(file_path: str) -> Dict[str, str]:
"""Load and preprocess data from a JSON file.
Parameters:
- file_path (str): The path to the JSON file to be loaded and processed.
Returns:
- Dict[str, str]: A dictionary containing the original keys from the JSON
file and their associated preprocessed text values.
If the file is not found, an empty dictionary is returned.
Raises:
- FileNotFoundError: If the specified file does not exist, a logging error
is recorded, and an empty dictionary is returned.
"""
try:
with open(file_path, mode="r", encoding="utf-8") as file:
data = json.load(file)
return {key: preprocess(text) for key, text in data.items()}
except FileNotFoundError:
logging.error("File %s not found.", file_path)
return {}
def generate_embeddings(texts: Dict[str, str]) -> Dict[str, List[float]]:
"""
Generate embeddings for a dictionary of texts.
Parameters:
- texts (Dict[str, str]): A dictionary where the keys are identifiers and the values are the texts for which embeddings will be generated.
Returns:
- Dict[str, List[float]]: A dictionary where the keys are the same identifiers from the input and the values are the generated embeddings as lists of floats.
"""
model = SentenceTransformer("bert-base-nli-mean-tokens")
return {key: model.encode([text])[0] for key, text in texts.items()} # type: ignore
def calculate_similarity(
tags_emb: Dict[str, List[float]], stereotypes_emb: Dict[str, List[float]]
) -> pd.DataFrame:
"""
Calculate cosine similarity between the two sets of embeddings.
Parameters:
- tags_emb (Dict[str, List[float]]): A dictionary where the keys are tag identifiers and the values are the embeddings of the tags.
- stereotypes_emb (Dict[str, List[float]]): A dictionary where the keys are stereotype identifiers and the values are the embeddings of the stereotypes.
Returns:
- pd.DataFrame: A DataFrame containing the cosine similarity scores between tags and stereotypes. The rows are indexed by tag identifiers, and the columns are labeled with stereotype identifiers.
"""
tags_keys, tags_values = zip(*tags_emb.items())
stereotypes_keys, stereotypes_values = zip(*stereotypes_emb.items())
sim_matrix = cosine_similarity(tags_values, stereotypes_values)
return pd.DataFrame(sim_matrix, index=tags_keys, columns=stereotypes_keys)
def plot_heatmap(sim_df: pd.DataFrame):
plt.figure(figsize=(10, 8))
sns.heatmap(sim_df, annot=True, cmap="coolwarm", fmt=".2f")
plt.title("Cosine Similarity between Videos and Stereotypes")
plt.xlabel("Stereotypes")
plt.ylabel("Videos")
plt.savefig(f"{PLOTS_PATH}video_stereotype_heatmap.png")
plt.show()
def create_network_graph(
tags: Dict[str, str], stereotypes: Dict[str, str], sim: pd.DataFrame
):
"""
Create and plot a network graph based on similarity scores between tags and stereotypes.
Parameters:
- tags (Dict[str, str]): A dictionary where keys are tag names and values are their descriptions.
- stereotypes (Dict[str, str]): A dictionary where keys are stereotype names and values are their descriptions.
- sim (pd.DataFrame): A DataFrame containing the cosine similarity scores between tags and stereotypes.
The function generates a network graph where nodes represent tags and stereotypes, and edges represent the similarity scores between them. Nodes are grouped by type (tag or stereotype), and edge thickness is proportional to the similarity score.
"""
G = nx.Graph()
G.add_nodes_from(tags.keys(), bipartite=0)
G.add_nodes_from(stereotypes.keys(), bipartite=1)
for i, tag in enumerate(tags.keys()):
for j, stereotype in enumerate(stereotypes.keys()):
weight = sim.iloc[i, j]
if weight > 0.3:
G.add_edge(tag, stereotype, weight=weight)
pos = nx.spring_layout(G, k=0.5, iterations=50)
plt.figure(figsize=(12, 8))
edges = G.edges(data=True)
weights = [edge[2]["weight"] for edge in edges]
nx.draw(
G,
pos,
with_labels=True,
node_size=3000,
node_color="skyblue",
font_size=10,
font_weight="bold",
)
nx.draw_networkx_edges(
G, pos, edge_color=weights, edge_cmap=plt.cm.Blues, width=2
)
plt.title(
"Network Graph of Cosine Similarity between Tags and Stereotypes"
)
plt.savefig(f"{PLOTS_PATH}video_stereotype_graph.png")
plt.show()
if __name__ == "__main__":
download_nltk_data()
stereotypes = load_and_preprocess_data(STEREOTYPES_JSON_PATH)
tags = load_and_preprocess_data(TAGS_JSON_PATH)
stereotypes_emb = generate_embeddings(stereotypes)
tags_emb = generate_embeddings(tags)
sim_df = calculate_similarity(tags_emb, stereotypes_emb)
plot_heatmap(sim_df)
create_network_graph(tags, stereotypes, sim_df)