app.py

import os
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

CACHE_DIR = "/app/cache"

# Ensure the cache directory is writable and exists
os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.makedirs(CACHE_DIR, exist_ok=True)  # Create the directory if it doesn't exist

# Load your Hugging Face model
model = pipeline("text-generation", model="devops-bda/Abap")

# Initialize FastAPI app
app = FastAPI()

# Define input format
class InputData(BaseModel):
    input_text: str

# Health check endpoint
@app.get("/health")
async def health_check():
    return {"status": "ok", "message": "Model is ready"}
    
# Define prediction endpoint
@app.post("/predict")
async def predict(data: InputData):
    result = model(data.input_text, max_length=500)
    return {"output": result}