Update main.py

main.py

@@ -1,36 +1,33 @@
-from fastapi import FastAPI, UploadFile, File
-from fastapi.staticfiles import StaticFiles
-import os
-
-app = FastAPI()
-
-# ربط مجلد static للواجهة الأمامية
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-@app.get("/")
-async def root():
-    return {"message": "Welcome to AI Web App"}
-
-from transformers import pipeline
-
-# تحميل نموذج لتلخيص النصوص
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-# تحميل نموذج لتفسير الصور (يمكن استبداله بـ Qwen إذا كان متاحًا)
-image_captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-
-from tika import parser
-
-@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"]}
-
+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"]}