from fastapi import FastAPI, UploadFile, File
from fastapi.staticfiles import StaticFiles
from transformers import pipeline
from tika import parser
import os

app = FastAPI()

# Mount static folder for frontend
app.mount("/static", StaticFiles(directory="static"), name="static")

# Define a writable cache directory inside the container
CACHE_DIR = "/app/cache"
os.makedirs(CACHE_DIR, exist_ok=True)  # Create the directory if it doesn’t exist

# Load models with custom cache directory
summarizer = pipeline("summarization", model="facebook/bart-large-cnn", cache_dir=CACHE_DIR)
image_captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", cache_dir=CACHE_DIR)

@app.get("/")
async def root():
    return {"message": "Welcome to AI Web App"}

@app.post("/summarize")
async def summarize_document(file: UploadFile = File(...)):
    content = parser.from_buffer(await file.read())["content"]
    summary = summarizer(content, max_length=130, min_length=30, do_sample=False)
    return {"summary": summary[0]["summary_text"]}

@app.post("/caption")
async def caption_image(file: UploadFile = File(...)):
    caption = image_captioner(await file.read())
    return {"caption": caption[0]["generated_text"]}