main.py

# -*- coding: utf-8 -*-
"""Prediction Server for Vertex AI

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1eml86CidWa8WEvCNAPPRwoFRmbNFVsuP
"""

from fastapi import FastAPI, Request
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import os

# Initialize the FastAPI app
app = FastAPI()

# --- Load the Fine-Tuned Model ---
# The model will be packaged inside the container with this script.
model_path = "./climatebert_misinfo_model"
classifier = None

# This event handler runs when the server starts up.
@app.on_event("startup")
def load_model():
    global classifier
    if os.path.isdir(model_path):
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSequenceClassification.from_pretrained(model_path)
        classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
        print("✅ ClimateBERT model loaded successfully!")
    else:
        print(f"--- FATAL ERROR: Model not found at {model_path} ---")

# Define the prediction endpoint
@app.post("/predict")
async def predict(request: Request):
    data = await request.json()
    text = data.get("text", "")

    if not text or not classifier:
        return {"error": "Text is missing or model is not loaded."}

    # The pipeline returns a list of dictionaries, e.g., [{'label': 'news', 'score': 0.99}]
    prediction = classifier(text)[0]
    return prediction

# Health check endpoint for Vertex AI
@app.get("/health")
def health():
    return {"status": "ok"}