114 lines
3.3 KiB
Python
114 lines
3.3 KiB
Python
import pandas as pd
|
|
from prefect import task, Flow, context
|
|
from pandas import DataFrame
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
|
|
from preprocessing import preprocess_text
|
|
|
|
|
|
@task
|
|
def get_train_set() -> DataFrame:
|
|
logger = context.get("logger")
|
|
train = pd.read_csv('train.csv')
|
|
train = train.drop(['keyword', 'location'], axis=1)
|
|
logger.info(f"Train set: {len(train)} elements")
|
|
return train
|
|
|
|
|
|
@task
|
|
def get_test_set() -> DataFrame:
|
|
logger = context.get("logger")
|
|
test = pd.read_csv('test.csv')
|
|
logger.info(f"Test set: {len(test)} elements")
|
|
return test
|
|
|
|
|
|
@task
|
|
def preprocess_train(train: DataFrame) -> DataFrame:
|
|
pp_text_train = []
|
|
for text_data in train['text']:
|
|
pp_text_data = preprocess_text(text_data)
|
|
pp_text_train.append(pp_text_data)
|
|
train['pp_text'] = pp_text_train
|
|
return train
|
|
|
|
|
|
@task
|
|
def preprocess_test(test: DataFrame) -> DataFrame:
|
|
pp_text_test = []
|
|
for text_data in test['text']:
|
|
pp_text_data = preprocess_text(text_data)
|
|
pp_text_test.append(pp_text_data)
|
|
test['pp_text'] = pp_text_test
|
|
return test
|
|
|
|
|
|
@task
|
|
def prepare_vectorizer(train_data: DataFrame, test_data: DataFrame) -> TfidfVectorizer:
|
|
train_text_data = list(train_data['pp_text'])
|
|
test_text_data = list(test_data['pp_text'])
|
|
corpus = train_text_data + test_text_data
|
|
tf = TfidfVectorizer()
|
|
fitted_vectorizer = tf.fit(corpus)
|
|
return fitted_vectorizer
|
|
|
|
|
|
@task
|
|
def transform_train(vectorizer: TfidfVectorizer, train_set: DataFrame) -> DataFrame:
|
|
return vectorizer.transform(train_set)
|
|
|
|
|
|
@task
|
|
def transform_test(vectorizer: TfidfVectorizer, test_set: DataFrame) -> DataFrame:
|
|
return vectorizer.transform(test_set)
|
|
|
|
|
|
@task
|
|
def split_test_set(X: DataFrame, Y: DataFrame) -> dict:
|
|
X_train, X_test, y_train, y_test = train_test_split(X, Y)
|
|
return {'X_train': X_train, 'X_test': X_test, 'y_train': y_train, 'y_test': y_test}
|
|
|
|
|
|
@task
|
|
def train_model(X: DataFrame, Y: DataFrame) -> LogisticRegression:
|
|
scikit_log_reg = LogisticRegression()
|
|
model = scikit_log_reg.fit(X, Y)
|
|
return model
|
|
|
|
|
|
@task
|
|
def evaluate(model: LogisticRegression, X: DataFrame, Y: DataFrame) -> None:
|
|
logger = context.get("logger")
|
|
predictions = model.predict(X)
|
|
count = 0
|
|
for guess, answer in zip(predictions, Y):
|
|
if guess == answer:
|
|
count += 1
|
|
score = count/len(Y)
|
|
logger.info(f"model score: {count/len(Y)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
with Flow("My First Prefect Flow!") as flow:
|
|
train_data = get_train_set()
|
|
test_data = get_test_set()
|
|
|
|
train_data = preprocess_train(train_data)
|
|
test_data = preprocess_test(test_data)
|
|
|
|
vectorizer = prepare_vectorizer(train_data, test_data)
|
|
|
|
vectorized_train_data = transform_train(vectorizer, train_data['pp_text'])
|
|
vectorized_test_data = transform_test(vectorizer, train_data['pp_text'])
|
|
|
|
splitted_data = split_test_set(vectorized_train_data, train_data['target'])
|
|
model = train_model(splitted_data['X_train'], splitted_data['y_train'])
|
|
evaluate(model, splitted_data['X_test'], splitted_data['y_test'])
|
|
|
|
flow.validate()
|
|
# flow.visualize()
|
|
flow.run()
|