Update main.py

main.py

@@ -9,7 +9,6 @@ from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
-from transformers import pipeline  # Pour le pipeline QA
 import fitz  # PyMuPDF
 from PIL import Image
 import io
@@ -21,22 +20,16 @@ from pptx import Presentation
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# 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
-Path(TEMP_CACHE_DIR).mkdir(parents=True, exist_ok=True)
-
 # Initialisation de l'application FastAPI
 app = FastAPI()
 
-# Configuration CORS avec méthodes explicites
+# Configuration CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_methods=["POST", "GET", "PUT", "DELETE", "OPTIONS"],  # Méthodes explicites
+    allow_methods=["POST", "GET", "PUT", "DELETE", "OPTIONS"],
     allow_headers=["*"],
-    allow_credentials=True,  # Permettre les credentials si nécessaire
+    allow_credentials=True,
 )
 
 # Chemins des fichiers
@@ -44,18 +37,15 @@ BASE_DIR = Path(__file__).parent
 UPLOAD_FOLDER = BASE_DIR / "uploads"
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
-# Modèles Hugging Face
+# Configuration des 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"  # Modèle QA léger
+    "qa": "istilbert-base-cased-distilled-squad"  # plus léger
 }
 
-# Pipeline QA pour distilbert-base-cased-distilled-squad
-qa_pipeline = pipeline("question-answering", model=MODELS["qa"], tokenizer=MODELS["qa"], cache_dir=TEMP_CACHE_DIR)
-
 # Modèles Pydantic
 class FileInfo(BaseModel):
     file_id: str
@@ -75,7 +65,7 @@ class QARequest(BaseModel):
     file_id: Optional[str] = None
     question: str
 
-# Fonctions utilitaires (extraction de texte, etc.)
+# Fonctions utilitaires
 def extract_text_from_pdf(file_path: str) -> str:
     try:
         doc = fitz.open(file_path)
@@ -122,11 +112,9 @@ async def process_uploaded_file(file: UploadFile) -> FileInfo:
     file_id = str(uuid.uuid4())
     file_path = str(UPLOAD_FOLDER / f"{file_id}{file_ext}")
     
-    # Sauvegarde du fichier
     with open(file_path, "wb") as buffer:
         buffer.write(await file.read())
     
-    # Extraction du texte selon le type de fichier
     text = ""
     if file_ext == ".pdf":
         text = extract_text_from_pdf(file_path)
@@ -145,12 +133,15 @@ async def process_uploaded_file(file: UploadFile) -> FileInfo:
         extracted_text=text if text else None
     )
 
-# Route de test pour vérifier l'API
+# Routes de l'API
 @app.get("/api/test")
 async def test_api():
     return {"status": "API working", "environment": "Hugging Face" if os.environ.get("HF_SPACE") else "Local"}
 
-# Endpoints (upload, summarize, caption, answer, etc.)
+@app.get("/api")
+async def api_root():
+    return {"status": "API is running"}
+
 @app.post("/api/upload")
 async def upload_files(files: List[UploadFile] = File(...)):
     logger.info(f"Upload request received with {len(files)} files")
@@ -165,6 +156,47 @@ 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[:5000],  # limite si le document est trop long
+            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:
@@ -182,22 +214,31 @@ async def answer_question(request: QARequest):
                 else:
                     with open(file_path, "r", encoding="utf-8") as f:
                         context = f.read()
-
-        # Utiliser le pipeline QA pour obtenir la réponse
-        result = qa_pipeline(question=request.question, context=context)
         
-        return {
-            "answer": result["answer"],
-            "confidence": result["score"]
-        }
-    except StopIteration:
-        logger.error("File not found")
-        raise HTTPException(404, "Fichier non trouvé")
+        if not context:
+            raise HTTPException(400, "Aucun contexte trouvé pour répondre à la question.")
+        
+        response = client.question_answering(
+            question=request.question,
+            context=context,
+            model=MODELS["qa"]
+        )
+        
+        return {"answer": response}
     except Exception as e:
         logger.error(f"QA error: {e}")
         raise HTTPException(500, f"Erreur de réponse: {str(e)}")
 
-# Gestion des erreurs
+@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 globales
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request, exc):
     return JSONResponse(
@@ -213,9 +254,9 @@ async def generic_exception_handler(request, exc):
         content={"detail": "Une erreur interne est survenue"},
     )
 
-# Montage des fichiers statiques APRÈS la définition des routes API
+# Montage des fichiers statiques
 app.mount("/", StaticFiles(directory=BASE_DIR, html=True), name="static")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)