tomy07417/Natural-Language-Processing-with-Disaster-Tweets

1import os2os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # CPU3 4import gradio as gr5import tensorflow as tf6from huggingface_hub import hf_hub_download7from transformers import AutoTokenizer, TFAutoModel8 9 10@tf.keras.utils.register_keras_serializable()11class DistilBertLayer(tf.keras.layers.Layer):12    def __init__(self, model_name="vinai/bertweet-base", **kwargs):13        super().__init__(**kwargs)14        self.model_name = model_name15        self.bert = TFAutoModel.from_pretrained(model_name, from_pt=True)16 17    def call(self, inputs):18        input_ids, attention_mask = inputs19        outputs = self.bert(20            input_ids=input_ids,21            attention_mask=attention_mask,22            training=False23        )24        return outputs.last_hidden_state25 26    def get_config(self):27        config = super().get_config()28        config.update({"model_name": self.model_name})29        return config30 31 32# 1) Repo donde subiste el .keras (MODELS, no Spaces)33MODEL_REPO = "tomy07417/disaster-tweets-bertweet-gru"  # <-- CAMBIÁ ESTO34MODEL_FILE = "bertweet_gru_model.keras"                # <-- nombre exacto en el repo35 36# 2) Descarga con cache (no lo baja cada vez)37model_path = hf_hub_download(38    repo_id=MODEL_REPO,39    filename=MODEL_FILE,40    repo_type="model"41)42 43# 3) Cargar el modelo desde el path descargado44model = tf.keras.models.load_model(45    model_path,46    custom_objects={"DistilBertLayer": DistilBertLayer},47    compile=False48)49 50tokenizer = AutoTokenizer.from_pretrained("vinai/bertweet-base")51 52 53def predict(text):54    inputs = tokenizer(55        [text],56        max_length=50,57        truncation=True,58        padding="max_length",59        return_tensors="tf"60    )61 62    input_ids = inputs["input_ids"]63    attention_mask = inputs["attention_mask"]64 65    # si tu salida es (1,) sigmoid:66    prob = model.predict([input_ids, attention_mask])[0][0]67    pred = bool(prob > 0.5)68 69    return {"prob": float(prob), "pred": pred}70 71 72demo = gr.Interface(73    fn=predict,74    inputs=gr.Textbox(lines=3, label="Tweet"),75    outputs=gr.JSON(label="Result"),76    title="Tweet classifier",77    description="Paste a tweet in English"78)79 80if __name__ == "__main__":81    # En Spaces NO uses share=True82    demo.launch(ssr_mode=False)83 84