sport-text-classification-b.../bayes2.py

34 lines
869 B
Python
Raw Normal View History

2021-05-12 19:58:28 +02:00
import pandas as pd
2021-05-12 20:18:46 +02:00
import numpy as np
2021-05-12 19:58:28 +02:00
import gzip
2021-05-12 20:29:37 +02:00
from sklearn.pipeline import make_pipeline
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer
2021-05-12 19:58:28 +02:00
2021-05-12 20:29:37 +02:00
# Read data
2021-05-12 19:58:28 +02:00
dev = pd.read_table('dev-0/in.tsv', error_bad_lines=False, header=None)
test = pd.read_table('test-A/in.tsv', error_bad_lines=False, header=None)
2021-05-12 20:18:46 +02:00
X_train = []
y_train = []
2021-05-12 19:58:28 +02:00
with gzip.open('train/train.tsv.gz', 'r') as f:
for l in f:
line = l.decode('UTF-8').replace("\n", "").split("\t")
2021-05-12 20:29:37 +02:00
y_train.append(int(line[0]))
X_train.append(str(line[1:]))
2021-05-12 19:58:28 +02:00
2021-05-12 20:29:37 +02:00
# Convert to unified types
X_train = np.asarray(X_train)
y_train = np.asarray(y_train)
2021-05-12 20:18:46 +02:00
X_dev = dev[0].values
X_test = test[0].values
2021-05-12 20:29:37 +02:00
print(type(y_train[0]))
print(X_train[0])
# Create model
model = make_pipeline(TfidfVectorizer(), MultinomialNB())
model.fit(X_train, y_train)