forked from kubapok/retroc2
110 lines
2.8 KiB
Plaintext
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
|
|
}
|