<jupyter_start><jupyter_text>Custom tokenisation<jupyter_code>train_ds = get_dataset("train")
valid_ds = get_dataset("valid")
test_ds = get_dataset("test")
vectorize_code = TextVectorization(
    standardize=None,
    max_tokens=vocab_size_code,
    output_mode="int",
    output_sequence_length=max_len_code,
)

vectorize_query = TextVectorization(
    max_tokens=vocab_size_query,
    output_mode="int",
    output_sequence_length=max_len_query,
)
vectorize_code.adapt(train_ds.map(lambda q,c: c))
vectorize_query.adapt(train_ds.map(lambda q,c: q))
def vectorize(query, code):
    query = tf.expand_dims(query, -1)
    code = tf.expand_dims(code, -1)
    return vectorize_query(query), vectorize_code(code)
train_ds = train_ds.map(vectorize)
valid_ds = valid_ds.map(vectorize)
test_ds = test_ds.map(vectorize)
train_ds = train_ds.prefetch(buffer_size=10)
valid_ds = valid_ds.prefetch(buffer_size=10)
test_ds = test_ds.prefetch(buffer_size=10)<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>loss_tracker = tf.keras.metrics.Mean(name="loss")
mrr_tracker = tf.keras.metrics.Mean(name="mrr")

class CustomModel(tf.keras.Model):
    def train_step(self, data):
        with tf.GradientTape() as tape:
            logits = self(data, training=True)
            loss_value = loss_(None, logits)
        
        mrr_value = mrr(None, logits)

        gradients = tape.gradient(loss_value, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        # Compute our own metrics
        loss_tracker.update_state(loss_value)
        mrr_tracker.update_state(mrr_value)
        return {"loss": loss_tracker.result(), "mrr": mrr_tracker.result()}
    
    def test_step(self, data):
            
        logits = self(data, training=False)
        loss_value = loss_(None, logits)
        mrr_value = mrr(None, logits)
        
        
        loss_tracker.update_state(loss_value)
        mrr_tracker.update_state(mrr_value)
        return {"loss": loss_tracker.result(), "mrr": mrr_tracker.result()}
input_query = tf.keras.Input(shape=(None,), dtype="int64")
input_q = Embedding(len(vectorize_query.get_vocabulary()) + 1, embeddings_dim_q)(input_query)

input_code = tf.keras.Input(shape=(None,), dtype="int64")
input_c = Embedding(len(vectorize_code.get_vocabulary()) + 1, embeddings_dim_c)(input_code)

lstm_q = LSTM(embeddings_dim_q, dropout=0.3, recurrent_dropout=0.0)(input_q)
lstm_c = LSTM(embeddings_dim_c, dropout=0.3, recurrent_dropout=0.0)(input_c)


model = CustomModel([input_query, input_code], [lstm_q, lstm_c])
def loss_(y_true, y_pred):
    
    q, c = y_pred

    similarity_score = tf.matmul(q, K.transpose(c))
    per_sample_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=similarity_score,
        labels=tf.range(q.shape[0])
    )
    return tf.reduce_sum(per_sample_loss) / q.shape[0]

def mrr(y_true, y_pred):
    q,c = y_pred
    
    similarity_scores = tf.matmul(q, K.transpose(c))
    
    # extract the logits from the diagonal of the matrix, which are the logits corresponding to the ground-truth
    correct_scores = tf.linalg.diag_part(similarity_scores)
    
    # compute how many queries have bigger logits than the ground truth (the diagonal) -> which will be incorrectly ranked
    compared_scores = similarity_scores >= tf.expand_dims(correct_scores, axis=-1)
    
    compared_scores = tf.cast(compared_scores, tf.float16)
    # for each row of the matrix (query), sum how many logits are larger than the ground truth
    # ...then take the reciprocal of that to get the MRR for each individual query (you will need to take the mean later)
    return K.mean(tf.math.reciprocal(tf.reduce_sum(compared_scores, axis=1)))
model.compile(optimizer='adam', run_eagerly=True)
checkpoint = tf.keras.callbacks.ModelCheckpoint("./exp/vectorise_model.h5", monitor='val_mrr', verbose=0, save_best_only=True, mode='max', save_freq='epoch')
tb = tf.keras.callbacks.TensorBoard(log_dir='./exp/vectorise_model1', update_freq='batch')
train_ds = train_ds.unbatch().batch(128)
model.fit(train_ds, validation_data=valid_ds, epochs=10, callbacks=[checkpoint, tb])<jupyter_output><empty_output><jupyter_text>Epoch 1/10  15/3221 [..............................] - ETA: 2:07:26 - loss: 4.8522 - mrr: 0.0409  Trainig using loop<jupyter_code>train_ds = train_ds.cache().prefetch(100)
valid_ds = valid_ds.cache().prefetch(100)
test_ds = valid_ds.cache().prefetch(100)
optimizer = tf.optimizers.Adam()


liveloss = PlotLosses()
logs = {}
epochs = 10

for epoch in range(epochs):

    # TRAINING
    losses = []
    mrrs = []
    for x in train_ds:
        
        with tf.GradientTape() as tape:
            logits = model(x)
            loss_value = loss_(None, logits)

        mrr_value = mrr(None, logits)

        # calculate gradient
        gradients = tape.gradient(loss_value, model.trainable_variables)

        # Update the weights
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))

        losses.append(loss_value)
        mrrs.append(mrr_value)

        print(f"Epoch: {epoch}; Loss: {loss_value}; MRR: {mrr_value} <- Train", end="\r")
        
    logs['loss'] = np.mean(losses)
    logs['mrr'] = np.mean(mrrs)
    
    
    # VALIDATION
    losses = []
    mrrs = []
    for x in valid_ds:
        
        logits = model(x)
        
        loss_value = loss_(None, logits)
        mrr_value = mrr(None, logits)
        
        losses.append(loss_value)
        mrrs.append(mrr_value)
        
        print(f"Epoch: {epoch}; Loss: {loss_value}; MRR: {mrr_value} <- Test", end="\r")
        
    logs['val_loss'] = np.mean(losses)
    logs['val_mrr'] = np.mean(mrrs)
    
    liveloss.update(logs)
    liveloss.send()<jupyter_output><empty_output><jupyter_text>Epoch: 0; Loss: 4.852148056030273; MRR: 0.0242462158203125 <- Trai  Fasttext Model<jupyter_code>embeddings_dim_c, embeddings_dim_q = 256, 256
max_len_code, max_len_query = 98, 48
vocab_size_code = 500000
vocab_size_query = 250000
batch_size = 8
# def _code_parse(row):
#     parser_ = get_parser(lang=row['language'], so_path="../resources/csnet_parse_build.so")
#     g_ = parse_program(row['code'], parser=parser_)
#     return g_

# model = ft.train_unsupervised(
#     "../resources/python_processed/query_corpus.txt",
#     lr = 0.005,
#     epoch = 30,
#     dim = 256,
#     thread = 40,
# )
# model.save_model("../resources/python_processed/query_ft.bin")

query_ft = ft.load_model("../resources/python_processed/query_ft.bin")
code_ft = ft.load_model("../resources/python_processed/code_no_ast.bin")
def get_dataset(split):
    df = get_data_df("../resources/data/", ["python"] , [split])
    
    q_s, c_s = [], []
    for i, row in df[['docstring_tokens','code_tokens']].iterrows():
        _query = [query_ft[token] for token in row[0]]
        _code = [code_ft[token] for token in row[1]]

        q_s.append(_query)
        c_s.append(_code)
    
    q_s = tf.keras.preprocessing.sequence.pad_sequences(
        q_s, maxlen=max_len_query, dtype='float', padding='post', truncating='post',
        value=[0]
    )

    c_s = tf.keras.preprocessing.sequence.pad_sequences(
        c_s, maxlen=max_len_code, dtype='float', padding='post', truncating='post',
        value=[0]
    )
    
    query_dataset = tf.data.Dataset.from_tensor_slices(q_s)
    code_dataset = tf.data.Dataset.from_tensor_slices(c_s)
    
    dataset = tf.data.Dataset.zip((query_dataset, code_dataset))
    
    return dataset.batch(batch_size)
train_ds = get_dataset("train")
valid_ds = get_dataset("valid")
# train_ds = train_ds.unbatch().batch(128)
def loss_(y_true, y_pred):
    
    q, c = y_pred

    similarity_score = tf.matmul(q, K.transpose(c))
    per_sample_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=similarity_score,
        labels=tf.range(q.shape[0])
    )
    return tf.reduce_sum(per_sample_loss) / q.shape[0]



def mrr(y_true, y_pred):
    q,c = y_pred
    
    similarity_scores = tf.matmul(q, K.transpose(c))
    
    # extract the logits from the diagonal of the matrix, which are the logits corresponding to the ground-truth
    correct_scores = tf.linalg.diag_part(similarity_scores)
    
    # compute how many queries have bigger logits than the ground truth (the diagonal) -> which will be incorrectly ranked
    compared_scores = similarity_scores >= tf.expand_dims(correct_scores, axis=-1)
    
    compared_scores = tf.cast(compared_scores, tf.float16)
    # for each row of the matrix (query), sum how many logits are larger than the ground truth
    # ...then take the reciprocal of that to get the MRR for each individual query (you will need to take the mean later)
    return K.mean(tf.math.reciprocal(tf.reduce_sum(compared_scores, axis=1)))


input_query = tf.keras.Input(shape=(max_len_query, embeddings_dim_q))
input_code = tf.keras.Input(shape=(max_len_code, embeddings_dim_c))

lstm_q = Bidirectional(LSTM(embeddings_dim_q, dropout=0.5, recurrent_dropout=0.0))(input_query)
lstm_c = Bidirectional(LSTM(embeddings_dim_c, dropout=0.5, recurrent_dropout=0.0))(input_code)

model = tf.keras.Model([input_query, input_code], [lstm_q, lstm_c])
# model.compile(optimizer='adam')
optimizer = tf.optimizers.Adam()


liveloss = PlotLosses()
logs = {}
epochs = 10
best_val_loss = 100000.0

for epoch in range(epochs):

    # TRAINING
    losses = []
    mrrs = []
    for x in train_ds:
        
        with tf.GradientTape() as tape:
            logits = model(x)
            loss_value = loss_(None, logits)

        mrr_value = mrr(None, logits)

        # calculate gradient
        gradients = tape.gradient(loss_value, model.trainable_variables)

        # Update the weights
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))

        losses.append(loss_value)
        mrrs.append(mrr_value)

        print(f"Epoch: {epoch}; Loss: {loss_value}; MRR: {mrr_value} <- Train", end="\r")
        
    logs['loss'] = np.mean(losses)
    logs['mrr'] = np.mean(mrrs)
    
    
    # VALIDATION
    losses = []
    mrrs = []
    for x in valid_ds:
        
        logits = model(x)
        
        loss_value = loss_(None, logits)
        mrr_value = mrr(None, logits)
        
        losses.append(loss_value)
        mrrs.append(mrr_value)
        
        print(f"Epoch: {epoch}; Loss: {loss_value}; MRR: {mrr_value} <- Test", end="\r")
        
    logs['val_loss'] = np.mean(losses)
    logs['val_mrr'] = np.mean(mrrs)
    
    if logs['val_loss'] < best_val_loss:
        best_val_loss = logs['val_loss']
        model.save("exp/code_no_ast_embedding_screath_try_bidirectional.h5")
    
    liveloss.update(logs)
    liveloss.send()
liveloss.send()
model = tf.keras.models.load_model("exp/code_no_ast_embedding_screath_try_bidirectional.h5")
test_ds = get_dataset("test")
    losses = []
    mrrs = []
    for x in get_dataset("test"):
        
        logits = model(x)
        
        loss_value = loss_(None, logits)
        mrr_value = mrr(None, logits)
        
        losses.append(loss_value)
        mrrs.append(mrr_value)
        
        print(f"Epoch: test; Loss: {loss_value}; MRR: {mrr_value} <- Test", end="\r")
np.mean(mrrs)<jupyter_output><empty_output><jupyter_text>Predit/Demo<jupyter_code>model = tf.keras.models.load_model("exp/sratch_keras_2.h5")
model.summary()
query_encoder = tf.keras.models.Model(model.get_layer('input_1').input, model.get_layer('lstm').output)
code_encoder = tf.keras.models.Model(model.get_layer('input_2').input, model.get_layer('lstm_1').output)
test_ds = get_dataset("test")
querys = []
codes = []

for q,c in test_ds:
    querys.extend(query_encoder(q))
    codes.extend(code_encoder(c))
from annoy import AnnoyIndex
import random

t = AnnoyIndex(256, 'angular')
for i in range(len(codes)):
    t.add_item(i, codes[i])
t.on_disk_build("exp/code_no_ast_embedding_screath_try.annoy_on_disk")
t.build(100)
from annoy import AnnoyIndex
t = AnnoyIndex(256, 'angular')
t.load('exp/code_no_ast_embedding_screath_try.annoyme')
t.save('exp/code_no_ast_embedding_screath_try.annoyme')
test_df = get_data_df("../resources/data/", ["python"] , ["test"])
result = t.get_nns_by_vector(codes[2155], n=2, include_distances=False)
result
correct = 0 
for i in range(len(querys)):
    result = t.get_nns_by_vector(querys[i], n=10, include_distances=False)
    if i in result:
        correct += 1
correct/len(querys)
with open("../resources/data/python_dedupe_definitions_v2.pkl", "rb") as f:
    import pickle 
    definations = pickle.load(f)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12916302360702647, 'fraction_numerical': 0.015499562832843176, 'mean_word_length': 3.6393805309734515, 'pattern_counts': {'":': 4, '<': 28, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '90812', 'license': 'MIT', 'n_tokens_mistral': 4731, 'n_tokens_neox': 4533, 'n_words': 798, 'path': 'notebook/ScratchTry.ipynb', 'repo_name': 'sedflix/CodeSearchNetExperiments', 'size': '12581'}