Spaces:
Runtime error
Runtime error
| 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) | |
| async def root(): | |
| return {"message": "Welcome to AI Web App"} | |
| 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"]} | |
| async def caption_image(file: UploadFile = File(...)): | |
| caption = image_captioner(await file.read()) | |
| return {"caption": caption[0]["generated_text"]} |