import tensorflow as tf
import os

print("Loading model...")
model = tf.keras.models.load_model("sentiment_model.h5")

# Convert with dynamic batch size (None instead of 256)
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 100], dtype=tf.float32)])
def predict(x):
    return model(x)

converter = tf.lite.TFLiteConverter.from_concrete_functions(
    [predict.get_concrete_function()]
)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

with open("sentiment_model.tflite", "wb") as f:
    f.write(tflite_model)

print(f"Size: {os.path.getsize('sentiment_model.tflite') / 1024 / 1024:.1f} MB")
print("Done! Dynamic batch size model created.")