retroc2/model.ipynb
2022-05-17 13:03:26 +02:00

110 lines
2.8 KiB
Plaintext

{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "8f5480f9-fa82-4150-acff-9309fdc43690",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Pipeline(steps=[('tfidfvectorizer', TfidfVectorizer()),\n",
" ('linearregression', LinearRegression())])"
]
},
"execution_count": 1,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import pandas as pd\n",
"from sklearn.feature_extraction.text import TfidfVectorizer\n",
"from sklearn.pipeline import make_pipeline\n",
"from sklearn.linear_model import LinearRegression\n",
"from sklearn.metrics import mean_squared_error\n",
"\n",
"with open('train/train.tsv', 'r', encoding='utf8') as file:\n",
" train_data = pd.read_csv(file, sep='\\t', names=['Begin', 'End', 'Title', 'Publisher', 'Text'])\n",
"\n",
" \n",
"X = train_data['Text']\n",
"Y = train_data['Begin']\n",
"\n",
"\n",
"model = make_pipeline(TfidfVectorizer(), LinearRegression())\n",
"model.fit(X, Y)\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "02e89f1c-a2d0-4d41-94a2-aa86b257069d",
"metadata": {},
"outputs": [],
"source": [
"def readFile(filename):\n",
" result = []\n",
" with open(filename, 'r', encoding=\"utf-8\") as file:\n",
" for line in file:\n",
" text = line.split(\"\\t\")[0].strip()\n",
" result.append(text)\n",
" return result\n",
"\n",
"def write_pred(filename, predictions):\n",
" with open(filename, \"w\") as file:\n",
" for pred in predictions:\n",
" file.write(str(pred) + \"\\n\")\n",
" \n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "b85f5e22-eafb-41ee-aa2c-20c338d42701",
"metadata": {},
"outputs": [],
"source": [
"\n",
"dev_0 = readFile('dev-0/in.tsv')\n",
"predict_dev_0 = model.predict(dev_0)\n",
"write_pred('dev-0/out.tsv', predict_dev_0)\n",
"\n",
"dev_1 = readFile('dev-1/in.tsv')\n",
"predict_dev_1 = model.predict(dev_1)\n",
"write_pred('dev-1/out.tsv', predict_dev_1)\n",
"\n",
"test_A = readFile('test-A/in.tsv')\n",
"predict_test_A = model.predict(test_A)\n",
"write_pred('test-A/out.tsv', predict_test_A)\n",
"\n",
"\n",
"\n",
"\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.7"
}
},
"nbformat": 4,
"nbformat_minor": 5
}