forked from kubapok/retroc2
50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
import pandas as pd
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.pipeline import make_pipeline
|
|
from sklearn.linear_model import LinearRegression
|
|
from sklearn.metrics import mean_squared_error
|
|
|
|
with open('train/train.tsv', 'r', encoding='utf8') as file:
|
|
train_data = pd.read_csv(file, sep='\t', names=['Begin', 'End', 'Title', 'Publisher', 'Text'])
|
|
|
|
|
|
def readFile(filename):
|
|
result = []
|
|
with open(filename, 'r', encoding="utf-8") as file:
|
|
for line in file:
|
|
text = line.split("\t")[0].strip()
|
|
result.append(text)
|
|
return result
|
|
|
|
def write_pred(filename, predictions):
|
|
with open(filename, "w") as file:
|
|
for pred in predictions:
|
|
file.write(str(pred) + "\n")
|
|
|
|
|
|
|
|
train_data = train_data[:10000]
|
|
|
|
X = train_data['Text']
|
|
Y = train_data['Begin']
|
|
|
|
|
|
model = make_pipeline(TfidfVectorizer(), LinearRegression())
|
|
model.fit(X, Y)
|
|
|
|
|
|
|
|
dev_0 = readFile('dev-0/in.tsv')
|
|
predict_dev_0 = model.predict(dev_0)
|
|
write_pred('dev-0/out.tsv', predict_dev_0)
|
|
|
|
dev_1 = readFile('dev-1/in.tsv')
|
|
predict_dev_1 = model.predict(dev_1)
|
|
write_pred('dev-1/out.tsv', predict_dev_1)
|
|
|
|
test_A = readFile('test-A/in.tsv')
|
|
predict_test_A = model.predict(test_A)
|
|
write_pred('test-A/out.tsv', predict_test_A)
|
|
|
|
|