abhinavGangwar/tiny-model-30m

1import tensorflow as tf2import numpy as np3import json4import os5from tensorflow.keras.preprocessing.text import tokenizer_from_json6from tensorflow.keras.preprocessing.sequence import pad_sequences7 8class EndpointHandler:9    def __init__(self, path=""):10        model_path = os.path.join(path, "tf_model")11        src_tokenizer_path = os.path.join(path, "src_tokenizer.json")12        trg_tokenizer_path = os.path.join(path, "trg_tokenizer.json")13        config_path = os.path.join(path, "config.json")14 15        with open(config_path, "r") as f:16            self.config = json.load(f)17        self.max_len = self.config.get("max_len", 40)18 19        with open(src_tokenizer_path, "r") as f:20            self.src_tokenizer = tokenizer_from_json(f.read())21        with open(trg_tokenizer_path, "r") as f:22            self.trg_tokenizer = tokenizer_from_json(f.read())23 24        self.start_token = self.trg_tokenizer.word_index.get("__start__")25        self.end_token = self.trg_tokenizer.word_index.get("__end__")26        self.index_word = {v: k for k, v in self.trg_tokenizer.word_index.items()}27 28        print("Loading model using tf.saved_model.load...")29        loaded_model = tf.saved_model.load(model_path)30        self.inference_func = loaded_model.signatures["serving_default"]31        32        print("✅ Model and tokenizers loaded successfully.")33 34    def predict(self, sentence: str) -> str:35        if not sentence or not isinstance(sentence, str):36            return "❌ Invalid input."37 38        src_seq = self.src_tokenizer.texts_to_sequences([sentence])39        src_padded = pad_sequences(src_seq, maxlen=self.max_len, padding="post")40        src_tensor = tf.convert_to_tensor(src_padded, dtype=tf.int32)41 42        decoder_input = tf.constant([[self.start_token]], dtype=tf.int32)43        result_tokens = []44 45        for i in range(self.max_len):46            seq_len = tf.shape(decoder_input)[1]47            padding_needed = self.max_len - seq_len48            paddings = tf.zeros([1, padding_needed], dtype=tf.int32)49            model_input_tgt = tf.concat([decoder_input, paddings], axis=1)50 51            predictions = self.inference_func(52                src=src_tensor, tgt_in=model_input_tgt53            )54            55            logits = predictions.get("output_0")56            next_id_tensor = tf.argmax(57                logits[:, seq_len - 1, :], axis=-1, output_type=tf.int3258            )59            60            # --- THIS IS THE CORRECTED LINE ---61            # Extract the scalar integer from the numpy array (e.g., array() -> 58)62            # --- THIS IS THE CORRECTED LINE ---63            # Extract the scalar integer from the numpy array (e.g., array() -> 58)64            next_id_scalar = next_id_tensor.numpy().item()65 66            if next_id_scalar == self.end_token:67                break68            69            result_tokens.append(next_id_scalar)70            next_id_reshaped = tf.reshape(next_id_tensor, [1, 1])71            decoder_input = tf.concat([decoder_input, next_id_reshaped], axis=-1)72 73        words = [self.index_word.get(tok, "<unk>") for tok in result_tokens]74        return " ".join(words)