Update main.py

main.py

@@ -9,6 +9,7 @@ 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
@@ -43,9 +44,12 @@ client = InferenceClient(token=HF_TOKEN)
 MODELS = {
     "summary": "facebook/bart-large-cnn",
     "caption": "Salesforce/blip-image-captioning-large",
-    "qa": "meta-llama/Llama-2-70b-chat-hf"
+    "qa": "distilbert-base-cased-distilled-squad"  # Nouveau modèle QA plus léger
 }
 
+# Pipeline QA pour distilbert-base-cased-distilled-squad
+qa_pipeline = pipeline("question-answering", model=MODELS["qa"], tokenizer=MODELS["qa"])
+
 # Modèles Pydantic
 class FileInfo(BaseModel):
     file_id: str
@@ -171,15 +175,6 @@ async def summarize_document(request: SummaryRequest):
             with open(file_path, "r", encoding="utf-8") as f:
                 text = f.read()
         
-        prompt = f"""
-        Résumez ce document de manière concise en français.
-        Concentrez-vous sur les points principaux.
-        Le résumé doit faire environ {request.max_length} mots.
-        
-        Document:
-        {text[:5000]}... [truncated]
-        """
-        
         summary = client.summarization(
             text=text,
             model=MODELS["summary"],
@@ -226,25 +221,17 @@ 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)
         
-        # Format du prompt adapté pour Llama 2
-        prompt = f"""<s>[INST] Vous êtes un assistant IA qui répond à des questions en français.
-        Répondez de manière précise et concise.
-        
-        Contexte: {context[:3000]}
-        
-        Question: {request.question} [/INST]"""
-        
-        # Utilisation de text_generation au lieu de chat_completion
-        response = client.text_generation(
-            prompt=prompt,
-            model=MODELS["qa"],
-            max_new_tokens=500,
-            temperature=0.7,
-            top_p=0.9
-        )
-        
-        return {"answer": response}
+        return {
+            "answer": result["answer"],
+            "confidence": result["score"]
+        }
+    except StopIteration:
+        logger.error("File not found")
+        raise HTTPException(404, "Fichier non trouvé")
     except Exception as e:
         logger.error(f"QA error: {e}")
         raise HTTPException(500, f"Erreur de réponse: {str(e)}")