import os
from transformers import BertTokenizer, BertConfig, TFBertModel
from fastapi import FastAPI

app = FastAPI()

MODEL_DIR = os.environ.get("MODEL_DIR", "/app/bert_tf")

# Guard: create dir if missing; avoid listing non-existent paths
os.makedirs(MODEL_DIR, exist_ok=True)

# Probe one level deep only if there are entries
candidates = [MODEL_DIR]
try:
    for x in os.listdir(MODEL_DIR):
        p = os.path.join(MODEL_DIR, x)
        if os.path.isdir(p):
            candidates.append(p)
except FileNotFoundError:
    pass

for d in candidates:
    if (os.path.isfile(os.path.join(d, "vocab.txt"))
        and os.path.isfile(os.path.join(d, "config.json"))):
        MODEL_DIR = d
        break

tok  = BertTokenizer(vocab_file=f"{MODEL_DIR}/vocab.txt", do_lower_case=True)
cfg  = BertConfig.from_json_file(f"{MODEL_DIR}/config.json")
model= TFBertModel.from_pretrained(MODEL_DIR, from_tf=True, config=cfg)