# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

import logging
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence
import os
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)


class InputSentences(object):
    def __init__(self, dirname):
        self.dirname = dirname
 
    def __iter__(self):
        for fname in os.listdir(self.dirname):
            for line in open(os.path.join(self.dirname, fname)):
                yield line.split()


# this is a very simple, use-case
sentences = [['Rome', 'Italy'], ['Beijing', 'China']]
# train word2vec on the two sentences
model = Word2Vec(sentences, min_count=1)

model.most_similar(positive=['Rome'], topn=1)

StarcoderdataJupyterScriptsDedupFiltered