Update main.py

main.py

@@ -38,17 +38,22 @@ BASE_DIR = Path(__file__).parent
 UPLOAD_FOLDER = BASE_DIR / "uploads"
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
+# Configuration du cache pour Hugging Face
+TEMP_CACHE_DIR = "/tmp/huggingface_cache"
+os.environ["TRANSFORMERS_CACHE"] = TEMP_CACHE_DIR
+os.environ["HF_HOME"] = TEMP_CACHE_DIR
+
 # Modèles Hugging Face
 HF_TOKEN = os.getenv("HF_TOKEN")
 client = InferenceClient(token=HF_TOKEN)
 MODELS = {
     "summary": "facebook/bart-large-cnn",
     "caption": "Salesforce/blip-image-captioning-large",
-    "qa": "distilbert-base-cased-distilled-squad"  # Nouveau modèle QA plus léger
+    "qa": "distilbert-base-cased-distilled-squad"  # Modèle QA léger
 }
 
 # Pipeline QA pour distilbert-base-cased-distilled-squad
-qa_pipeline = pipeline("question-answering", model=MODELS["qa"], tokenizer=MODELS["qa"])
+qa_pipeline = pipeline("question-answering", model=MODELS["qa"], tokenizer=MODELS["qa"], cache_dir=TEMP_CACHE_DIR)
 
 # Modèles Pydantic
 class FileInfo(BaseModel):
@@ -69,7 +74,7 @@ class QARequest(BaseModel):
     file_id: Optional[str] = None
     question: str
 
-# Fonctions utilitaires
+# Fonctions utilitaires (extraction de texte, etc.)
 def extract_text_from_pdf(file_path: str) -> str:
     try:
         doc = fitz.open(file_path)
@@ -144,11 +149,7 @@ async def process_uploaded_file(file: UploadFile) -> FileInfo:
 async def test_api():
     return {"status": "API working", "environment": "Hugging Face" if os.environ.get("HF_SPACE") else "Local"}
 
-# Endpoints
-@app.get("/api")
-async def api_root():
-    return {"status": "API is running"}
-
+# Endpoints (upload, summarize, caption, answer, etc.)
 @app.post("/api/upload")
 async def upload_files(files: List[UploadFile] = File(...)):
     logger.info(f"Upload request received with {len(files)} files")
@@ -163,47 +164,6 @@ async def upload_files(files: List[UploadFile] = File(...)):
         logger.error(f"Upload error: {e}")
         raise HTTPException(500, f"Erreur lors de l'upload: {str(e)}")
 
-@app.post("/api/summarize")
-async def summarize_document(request: SummaryRequest):
-    try:
-        file_path = next(f for f in UPLOAD_FOLDER.glob(f"{request.file_id}*"))
-        text = ""
-        
-        if file_path.suffix == ".pdf":
-            text = extract_text_from_pdf(str(file_path))
-        else:
-            with open(file_path, "r", encoding="utf-8") as f:
-                text = f.read()
-        
-        summary = client.summarization(
-            text=text,
-            model=MODELS["summary"],
-            parameters={"max_length": request.max_length}
-        )
-        
-        return {"summary": summary}
-    except Exception as e:
-        logger.error(f"Summarization error: {e}")
-        raise HTTPException(500, f"Erreur de résumé: {str(e)}")
-
-@app.post("/api/caption")
-async def caption_image(request: CaptionRequest):
-    try:
-        file_path = next(f for f in UPLOAD_FOLDER.glob(f"{request.file_id}*"))
-        
-        with open(file_path, "rb") as image_file:
-            image_data = image_file.read()
-        
-        caption = client.image_to_text(
-            image=image_data,
-            model=MODELS["caption"]
-        )
-        
-        return {"caption": caption}
-    except Exception as e:
-        logger.error(f"Captioning error: {e}")
-        raise HTTPException(500, f"Erreur de description: {str(e)}")
-
 @app.post("/api/answer")
 async def answer_question(request: QARequest):
     try:
@@ -236,15 +196,6 @@ async def answer_question(request: QARequest):
         logger.error(f"QA error: {e}")
         raise HTTPException(500, f"Erreur de réponse: {str(e)}")
 
-@app.get("/api/file/{file_id}")
-async def get_file(file_id: str):
-    try:
-        file_path = next(f for f in UPLOAD_FOLDER.glob(f"{file_id}*"))
-        return FileResponse(file_path)
-    except Exception as e:
-        logger.error(f"File retrieval error: {e}")
-        raise HTTPException(404, "Fichier non trouvé")
-
 # Gestion des erreurs
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request, exc):