38 lines
869 B
Python
38 lines
869 B
Python
|
import pandas as pd
|
||
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
||
|
|
from sklearn.linear_model import LinearRegression
|
||
|
|
from sklearn.pipeline import make_pipeline
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def load_data(path):
|
||
|
|
return pd.read_csv(path, sep='\t', header=None)
|
||
|
|
|
||
|
|
|
||
|
|
def write_res(data, path):
|
||
|
|
with open(path, 'w') as f:
|
||
|
|
for line in data:
|
||
|
|
f.write(f'{line}\n')
|
||
|
|
print(f"Data written {path}/out.tsv")
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
train = load_data('train/train.tsv')
|
||
|
|
text = train[4]
|
||
|
|
year = train.apply(lambda row: ((row[0] + row[1])/2), axis=1)
|
||
|
|
|
||
|
|
model = make_pipeline(TfidfVectorizer(), LinearRegression())
|
||
|
|
model.fit(text, year)
|
||
|
|
|
||
|
|
|
||
|
|
for path in ['dev-0', 'dev-1', 'test-A']:
|
||
|
|
in_df = load_data(f'{path}/in.tsv')
|
||
|
|
predict = model.predict(in_df[0])
|
||
|
|
write_res(predict, f'{path}/out.tsv')
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|