polish-urban-legends-public/solution-dev0.py
2021-04-15 18:16:41 +02:00

24 lines
683 B
Python

from sklearn.feature_extraction.text import TfidfVectorizer
import numpy as np
import sklearn.metrics
from sklearn.cluster import KMeans
def main():
with open("dev-0/in.tsv") as in_file:
documents = in_file.readlines()
vectorizer = TfidfVectorizer(ngram_range=(1, 3), use_idf=False)
vectorizer = TfidfVectorizer()
document_vectors = vectorizer.fit_transform(documents)
predictions = KMeans(n_clusters=45).fit_predict(document_vectors)
with open("dev-0/out.tsv", "w") as out_file:
for prediction in predictions:
out_file.write(str(prediction) + '\n')
if __name__ == '__main__':
main()