151 lines
6.4 KiB
Python
151 lines
6.4 KiB
Python
import argparse
|
|
import os
|
|
from typing import List
|
|
|
|
import tensorflow as tf
|
|
from sklearn.pipeline import Pipeline
|
|
from tensorflow.keras.preprocessing.sequence import pad_sequences
|
|
from tensorflow.keras.preprocessing.text import Tokenizer
|
|
|
|
from misc import GENDER_MODELS_DIR, load_pickle
|
|
|
|
|
|
def predict_logreg(names: List[str], threshold: float):
|
|
"""
|
|
Predict gender labels for given names using a logistic regression model.
|
|
|
|
The function takes in a list of names and predicts the gender labels
|
|
based on a logistic regression model. A probabilistic threshold is used
|
|
to classify the names into one of the defined labels.
|
|
|
|
:param names:
|
|
A list of names for which the gender needs to be predicted. Each
|
|
name must be a string.
|
|
:param threshold:
|
|
A float value representing the threshold for classification. Names
|
|
with predicted probabilities greater than or equal to this value
|
|
will be classified into the positive class.
|
|
:return:
|
|
A tuple containing the predicted gender labels and their
|
|
corresponding probabilities. The first element of the tuple is a
|
|
list of predicted labels, while the second element is an array of
|
|
probability scores for each label.
|
|
"""
|
|
model_path = os.path.join(GENDER_MODELS_DIR, "regression_model.pkl")
|
|
encoder_path = os.path.join(GENDER_MODELS_DIR, "regression_label_encoder.pkl")
|
|
|
|
model: Pipeline = load_pickle(model_path)
|
|
label_encoder = load_pickle(encoder_path)
|
|
|
|
X = [name.lower().strip() for name in names]
|
|
proba = model.predict_proba(X)
|
|
pred = (proba[:, 1] >= threshold).astype(int)
|
|
labels = label_encoder.inverse_transform(pred)
|
|
return labels, proba
|
|
|
|
|
|
def predict_lstm(names: List[str], threshold: float, max_len=6):
|
|
"""
|
|
Predicts gender labels and probabilities for a list of names using a pre-trained BiLSTM model.
|
|
|
|
The function loads the model, tokenizer, and label encoder, performs preprocessing on the input
|
|
names, and then uses the loaded model to predict gender probabilities. Based on the threshold
|
|
value, it determines the predicted gender labels.
|
|
|
|
:param names: List of names to be classified.
|
|
:type names: List[str]
|
|
:param threshold: Probability threshold for classifying gender. If the predicted probability for the
|
|
'positive' class is greater than or equal to this threshold, it is classified accordingly.
|
|
:type threshold: float
|
|
:param max_len: Maximum length for name sequences. Names longer than this will be truncated, and shorter
|
|
ones will be padded. Default value is 6.
|
|
:type max_len: int, optional
|
|
|
|
:return: A tuple containing predicted labels and associated probabilities. Labels are the predicted gender
|
|
categories, and probabilities are the prediction scores for each input name.
|
|
:rtype: Tuple[numpy.ndarray, numpy.ndarray]
|
|
"""
|
|
model_path = os.path.join(GENDER_MODELS_DIR, "BiLSTM_model.h5")
|
|
tokenizer_path = os.path.join(GENDER_MODELS_DIR, "BiLSTM_tokenizer.pkl")
|
|
encoder_path = os.path.join(GENDER_MODELS_DIR, "BiLSTM_label_encoder.pkl")
|
|
|
|
model = tf.keras.models.load_model(model_path)
|
|
tokenizer: Tokenizer = load_pickle(tokenizer_path)
|
|
label_encoder = load_pickle(encoder_path)
|
|
|
|
X = tokenizer.texts_to_sequences([n.lower().strip() for n in names])
|
|
X = pad_sequences(X, maxlen=max_len, padding="post")
|
|
proba = model.predict(X)
|
|
pred = (proba[:, 1] >= threshold).astype(int)
|
|
labels = label_encoder.inverse_transform(pred)
|
|
return labels, proba
|
|
|
|
|
|
def predict_transformer(names: List[str], threshold: float, max_len=6):
|
|
"""
|
|
Predicts gender labels for the provided names using a pre-trained transformer model.
|
|
|
|
This function loads a pre-trained transformer model along with its tokenizer and label
|
|
encoder, converts input names into tokenized sequences, and processes them to generate
|
|
gender predictions. The function returns the predicted labels and the associated
|
|
probabilities for each sample.
|
|
|
|
:param names: List of names to predict gender labels for.
|
|
:type names: List[str]
|
|
:param threshold: Threshold value to determine the prediction class. Probability values
|
|
above or equal to the threshold will be assigned to one class, and those below to
|
|
another.
|
|
:type threshold: float
|
|
:param max_len: Maximum length for the sequences. Names will be truncated or padded to
|
|
this length during processing, default is 6.
|
|
:type max_len: int, optional
|
|
:return: A tuple containing two elements: a list of predicted gender labels as strings
|
|
and a NumPy array of probabilities for each gender class (where the first index
|
|
corresponds to one class, and the second index corresponds to another).
|
|
:rtype: Tuple[List[str], numpy.ndarray]
|
|
"""
|
|
model_path = os.path.join(GENDER_MODELS_DIR, "transformer.h5")
|
|
tokenizer_path = os.path.join(GENDER_MODELS_DIR, "transformer_tokenizer.pkl")
|
|
encoder_path = os.path.join(GENDER_MODELS_DIR, "transformer_label_encoder.pkl")
|
|
|
|
model = tf.keras.models.load_model(model_path)
|
|
tokenizer: Tokenizer = load_pickle(tokenizer_path)
|
|
label_encoder = load_pickle(encoder_path)
|
|
|
|
X = tokenizer.texts_to_sequences([n.lower().strip() for n in names])
|
|
X = pad_sequences(X, maxlen=max_len, padding="post")
|
|
proba = model.predict(X)
|
|
pred = (proba[:, 1] >= threshold).astype(int)
|
|
labels = label_encoder.inverse_transform(pred)
|
|
return labels, proba
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Predict gender from names using trained model")
|
|
parser.add_argument("--model", choices=["logreg", "lstm", "transformer"], required=True)
|
|
parser.add_argument("--name", nargs="+", required=True, help="One or more names")
|
|
parser.add_argument("--threshold", type=float, default=0.5, help="Threshold for classification")
|
|
args = parser.parse_args()
|
|
|
|
model = args.model
|
|
names = args.name
|
|
threshold = args.threshold
|
|
|
|
if model == "logreg":
|
|
labels, proba = predict_logreg(names, threshold)
|
|
elif model == "lstm":
|
|
labels, proba = predict_lstm(names, threshold)
|
|
elif model == "transformer":
|
|
labels, proba = predict_transformer(names, threshold)
|
|
else:
|
|
raise ValueError(f"Unsupported model type: {model}")
|
|
|
|
for i, name in enumerate(names):
|
|
p_female = proba[i][0]
|
|
p_male = proba[i][1]
|
|
print(f"{name} → {labels[i]} | P(f): {p_female:.2f} | P(m): {p_male:.2f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|