retroc2/run.py

38 lines
869 B
Python
Raw Normal View History

2022-06-07 22:19:51 +02:00
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()