58 lines
1.3 KiB
Python
58 lines
1.3 KiB
Python
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.linear_model import LinearRegression
|
|
import lzma
|
|
|
|
|
|
def generate_out(folder_path):
|
|
print('Generating out')
|
|
X_dev = []
|
|
Y_dev = []
|
|
with open(f'{folder_path}/in.tsv', 'r') as file:
|
|
for line in file:
|
|
line = line.strip()
|
|
X_dev.append(line)
|
|
|
|
print("step 5")
|
|
X_dev = vectorizer.transform(X_dev)
|
|
prediction = model.predict(X_dev)
|
|
|
|
print("step 6")
|
|
f = open(f"{folder_path}/out.tsv", "a")
|
|
for p in prediction:
|
|
f.write(str(p) + '\n')
|
|
f.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
X = []
|
|
Y = []
|
|
with lzma.open('train/train.tsv.xz', 'r') as file:
|
|
for line in file:
|
|
line = line.strip()
|
|
X.append(line.decode("utf-8"))
|
|
print("step 1")
|
|
with lzma.open('train/meta.tsv.xz', 'r') as file:
|
|
for line in file:
|
|
line = line.strip()
|
|
line = line.decode("utf-8")
|
|
Y.append(int(line.split('\t')[5]))
|
|
print("step 2")
|
|
vectorizer = TfidfVectorizer()
|
|
X = vectorizer.fit_transform(X)
|
|
|
|
print("step 3")
|
|
model = LinearRegression()
|
|
model.fit(X, Y)
|
|
|
|
print("step 4")
|
|
generate_out('dev-0')
|
|
generate_out('dev-1')
|
|
generate_out('test-A')
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|