CBo W-TF

from tensorflow import keras
import tensorflow as tf
import numpy as np

# course-edition deterministic TensorFlow news fixture v1
_course_titles = tf.constant([
    "Regional leaders agree on a public health plan",
    "The home team wins the basketball final",
    "Markets rise after a strong business report",
    "Researchers release an efficient intelligence model",
    "Diplomats begin a new round of peace talks",
    "The tennis champion returns after an injury",
    "Retail sales improve with consumer confidence",
    "A software update improves privacy",
])
_course_descriptions = tf.constant([
    "World policy and international cooperation",
    "Sports results and player development",
    "Business revenue and economic indicators",
    "Science and technology research",
    "World diplomacy and humanitarian aid",
    "Sports training and competition",
    "Business investment and employment",
    "Science security and technology",
])
_course_labels = tf.constant([0, 1, 2, 3, 0, 1, 2, 3], dtype=tf.int64)
_course_records = {"title": _course_titles, "description": _course_descriptions, "label": _course_labels}
ds_train = tf.data.Dataset.from_tensor_slices(_course_records).take(6)
ds_test = tf.data.Dataset.from_tensor_slices(_course_records).skip(6)

vocab_size = 5000

vectorizer = keras.layers.TextVectorization(
    max_tokens=vocab_size,
    output_mode="int",
)
embedder = keras.layers.Embedding(input_dim=vocab_size, output_dim=30)

model = keras.Sequential([
    keras.Input(shape=(1,), dtype=tf.int64),
    embedder,
    keras.layers.Dense(vocab_size, activation="softmax"),
])

model.summary()

def extract_text(x):
    return x['title']+' '+x['description']

vectorizer.adapt(ds_train.take(500).map(extract_text))
vocab = vectorizer.get_vocabulary()

def to_cbow(sent,window_size=2):
    res = []
    for i,x in enumerate(sent):
        for j in range(max(0,i-window_size),min(i+window_size+1,len(sent))):
            if i!=j:
                res.append([sent[j],x])
    return res

print(to_cbow(['I','like','to','train','networks']))
print(to_cbow(vectorizer('I like to train networks')))

X = []
Y = []
for i,x in zip(range(10000),ds_train.map(extract_text).as_numpy_iterator()):
    for w1, w2 in to_cbow(vectorizer(x),window_size=1):
        X.append(tf.expand_dims(w1,0))
        Y.append(tf.expand_dims(w2,0))

ds = tf.data.Dataset.from_tensor_slices((X,Y)).batch(256)

model.compile(optimizer=keras.optimizers.SGD(learning_rate=0.1),loss='sparse_categorical_crossentropy')
model.fit(ds, epochs=10)

vectors = embedder(vectorizer(vocab))
vectors = tf.reshape(vectors,(-1,30)) # we need reshape to get rid of extra dimension

paris_vec = embedder(vectorizer('paris'))[0]
print(paris_vec)

def close_words(x,n=5):
  vec = embedder(vectorizer(x))[0]
  top5 = np.linalg.norm(vectors-vec,axis=1).argsort()[:n]
  return [ vocab[x] for x in top5 ]

close_words('paris')

close_words('china')

close_words('official')