app.py

import os
import json
import pickle
import re
import numpy as np

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.preprocessing.sequence import pad_sequences
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel




app = FastAPI(title="PlantField Chatbot API", description="API untuk inference chatbot pertanian menggunakan model LSTM Seq2Seq.")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Pydantic Models ──
class ChatRequest(BaseModel):
    pertanyaan: str

class ChatResponse(BaseModel):
    jawaban: str

# ── Custom Layers (Dibutuhkan untuk memuat model Keras ──
class BahdanauAttention(layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.W1 = layers.Dense(units)
        self.W2 = layers.Dense(units)
        self.V  = layers.Dense(1)

    def call(self, query, values):
        query_exp = tf.expand_dims(query, 1)
        score = self.V(tf.nn.tanh(self.W1(values) + self.W2(query_exp)))
        weights  = tf.nn.softmax(score, axis=1)
        context  = tf.reduce_sum(weights * values, axis=1)
        return context, tf.squeeze(weights, -1)

    def get_config(self):
        config = super().get_config()
        return config

class AttentionDecoder(layers.Layer):
    def __init__(self, vocab_size, lstm_units, dropout, max_output, **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.lstm_units = lstm_units
        self.dropout_rate = dropout
        self.max_output = max_output
        
        self.attention = BahdanauAttention(lstm_units)
        self.lstm = layers.LSTM(
            lstm_units,
            return_sequences=True,
            return_state=True,
            dropout=dropout
        )
        self.concat = layers.Concatenate(axis=-1)
        self.layernorm = layers.LayerNormalization()
        self.dropout_layer = layers.Dropout(dropout)
        self.dense = layers.Dense(vocab_size, activation='softmax')

    def call(self, dec_emb, enc_out, state_h, state_c):
        dec_outputs = []

        for t in range(self.max_output):
            dec_tok = dec_emb[:, t:t+1, :]
            context, _ = self.attention(state_h, enc_out)
            context_exp = tf.expand_dims(context, 1)

            dec_in_combined = self.concat([dec_tok, context_exp])

            dec_out, state_h, state_c = self.lstm(
                dec_in_combined,
                initial_state=[state_h, state_c]
            )
            dec_outputs.append(dec_out)

        dec_outputs = tf.concat(dec_outputs, axis=1)
        dec_outputs = self.layernorm(dec_outputs)
        dec_outputs = self.dropout_layer(dec_outputs)
        output = self.dense(dec_outputs)
        return output

    def get_config(self):
        config = super().get_config()
        config.update({
            "vocab_size": self.vocab_size,
            "lstm_units": self.lstm_units,
            "dropout": self.dropout_rate,
            "max_output": self.max_output
        })
        return config

def masked_loss(y_true, y_pred):
    loss_fn = keras.losses.SparseCategoricalCrossentropy(reduction='none')
    loss    = loss_fn(y_true, y_pred)
    mask    = tf.cast(tf.not_equal(y_true, 0), dtype=loss.dtype)
    loss    = loss * mask
    return tf.reduce_sum(loss) / tf.reduce_sum(mask)

def masked_accuracy(y_true, y_pred):
    pred  = tf.cast(tf.argmax(y_pred, axis=-1), tf.int32)
    true  = tf.cast(y_true, tf.int32)
    match = tf.cast(tf.equal(pred, true), tf.float32)
    mask  = tf.cast(tf.not_equal(true, 0), tf.float32)
    return tf.reduce_sum(match * mask) / tf.reduce_sum(mask)


def build_seq2seq(vocab_size, embed_dim, lstm_units, max_input, max_output, dropout=0.3):
    embedding = layers.Embedding(vocab_size, embed_dim, mask_zero=True, name='shared_embedding')
    
    enc_input = layers.Input(shape=(max_input,), name='encoder_input')
    enc_emb   = embedding(enc_input)
    enc_emb   = layers.Dropout(dropout)(enc_emb)
    
    enc_out, fwd_h, fwd_c, bwd_h, bwd_c = layers.Bidirectional(
        layers.LSTM(lstm_units, return_sequences=True, return_state=True, dropout=dropout),
        name='encoder_bilstm'
    )(enc_emb)
    
    enc_h = layers.Concatenate()([fwd_h, bwd_h])
    enc_c = layers.Concatenate()([fwd_c, bwd_c])
    
    dec_lstm_units = lstm_units * 2
    
    dec_input = layers.Input(shape=(max_output,), name='decoder_input')
    dec_emb   = embedding(dec_input)
    dec_emb   = layers.Dropout(dropout)(dec_emb)
    
    decoder_layer = AttentionDecoder(
        vocab_size=vocab_size, lstm_units=dec_lstm_units, dropout=dropout,
        max_output=max_output, name='attention_decoder'
    )
    
    output = decoder_layer(dec_emb, enc_out, enc_h, enc_c)
    
    model = keras.models.Model(inputs=[enc_input, dec_input], outputs=output, name='PlantField_Seq2Seq')
    return model

# ── Global Objects ──
model = None
tokenizer = None
index_word = None
CONFIG = {}

# ── Helpers ──
def clean_text(text: str) -> str:
    text = text.lower().strip()
    text = re.sub(r"[^a-z0-9\s\?\.,'-]", ' ', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def decode_sequence(input_seq, _model, _tokenizer, idx_word, bos_idx, eos_idx, max_output_len):
    result_tokens = []
    
    dec_input_full = np.zeros((1, max_output_len), dtype=np.int32)
    dec_input_full[0, 0] = bos_idx

    for t in range(1, max_output_len):
        pred = _model.predict([input_seq, dec_input_full], verbose=0)
        token_id = np.argmax(pred[0, t-1, :])

        if token_id == eos_idx or token_id == 0:
            break

        result_tokens.append(token_id)
        dec_input_full[0, t] = token_id

    words = [idx_word.get(t, '') for t in result_tokens if idx_word.get(t, '')]
    return ' '.join(words)

# ── API Startup ──
@app.on_event("startup")
async def load_assets():
    global model, tokenizer, index_word, CONFIG
    
    # 1. Load Config
    try:
        with open("inference_config.json", "r") as f:
            CONFIG = json.load(f)
        print("Config ter-load dengan baik.")
    except Exception as e:
        print(f"Error loading config: {e}")
        
    # 2. Load Tokenizer
    try:
        with open("tokenizer.pkl", "rb") as f:
            tokenizer = pickle.load(f)
        index_word = {v: k for k, v in tokenizer.word_index.items()}
        print("Tokenizer ter-load dengan baik.")
    except Exception as e:
        print(f"Error loading tokenizer: {e}")

    # 3. Load Model
    try:
        import zipfile
        import os
        if not os.path.exists("/tmp/model.weights.h5"):
            with zipfile.ZipFile("plantfield_seq2seq.keras", "r") as z:
                z.extract("model.weights.h5", "/tmp")

        model = build_seq2seq(
            vocab_size = CONFIG.get('VOCAB_SIZE'),
            embed_dim  = CONFIG.get('EMBED_DIM'),
            lstm_units = CONFIG.get('LSTM_UNITS'),
            max_input  = CONFIG.get('MAX_INPUT_LEN'),
            max_output = CONFIG.get('MAX_OUTPUT_LEN'),
            dropout    = 0.3
        )
        model.load_weights("/tmp/model.weights.h5")
        print("Model ter-load dengan baik.")
    except Exception as e:
        print(f"Error loading model: {e}")


# ── Endpoint ──
@app.get("/")
def home():
    return {"status": "ok", "message": "API PlantField Chatbot berjalan dengan lancar."}

@app.post("/predict", response_model=ChatResponse)
async def predict_chatbot(req: ChatRequest):
    if not model or not tokenizer:
        raise HTTPException(status_code=500, detail="Model atau tokenizer belum ter-load.")

    pertanyaan = req.pertanyaan
    if not pertanyaan.strip():
        raise HTTPException(status_code=400, detail="Pertanyaan tidak boleh kosong.")

    # 1. Preprocessing
    cleaned = clean_text(pertanyaan)
    seq = tokenizer.texts_to_sequences([cleaned])
    
    if len(seq[0]) == 0:
        return ChatResponse(jawaban="Maaf, saya tidak mengerti maksud Anda. Bisa diperjelas?")

    # 2. Padding
    padded = pad_sequences(
        seq, 
        maxlen=CONFIG.get('MAX_INPUT_LEN', 25), 
        padding='post', 
        truncating='post'
    )

    # 3. Decoding
    jawaban = decode_sequence(
        padded, 
        model, 
        tokenizer, 
        index_word,
        CONFIG.get('BOS_IDX', 7),
        CONFIG.get('EOS_IDX', 8),
        CONFIG.get('MAX_OUTPUT_LEN', 35)
    )

    return ChatResponse(jawaban=jawaban)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)