Update main.py

main.py

@@ -74,42 +74,260 @@ def get_qa_model():
 #########################################################
 
 
-models_cache: Dict[str, pipeline] = {}
-
-def get_model(model_name: str, task: str):
-    if model_name not in models_cache:
-        models_cache[model_name] = pipeline(task, model=model_name)
-    return models_cache[model_name]
+# CORS Configuration
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
+# Constants
+MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
+MAX_TEXT_LENGTH = 2000
+MAX_QUESTION_LENGTH = 500
+MIN_QUESTION_LENGTH = 3
+SUPPORTED_LANGUAGES = {"fr", "en", "es", "de"}
+DEFAULT_LANGUAGE = "fr"
+
+SUPPORTED_FILE_TYPES = {
+    "docx": "Word Document",
+    "xlsx": "Excel Spreadsheet",
+    "pptx": "PowerPoint Presentation",
+    "pdf": "PDF Document",
+    "jpg": "JPEG Image",
+    "jpeg": "JPEG Image",
+    "png": "PNG Image"
+}
+
+
+
+
+
+MODEL_MAPPING = {
+    "fr": {
+        "qa": "illuin/camembert-base-fquad",
+        "summarization": "moussaKam/barthez-orangesum-abstract",
+        "translation": "Helsinki-NLP/opus-mt-fr-en"
+    },
+    "en": {
+        "qa": "deepset/roberta-base-squad2",
+        "summarization": "facebook/bart-large-cnn",
+        "translation": "Helsinki-NLP/opus-mt-en-fr"
+    },
+    "default": {
+        "image_captioning": "Salesforce/blip-image-captioning-large",
+        "multilingual_translation": "facebook/nllb-200-distilled-600M"
+    }
+}
+
+# Models cache
+models_cache: Dict[str, Pipeline] = {}
+
+# Pydantic Models
+class TranslationRequest(BaseModel):
+    text: constr(min_length=1, max_length=5000)
+    target_lang: constr(min_length=2, max_length=5)
+    src_lang: Optional[constr(min_length=2, max_length=5)] = None
+
+    @validator('target_lang', 'src_lang')
+    def validate_language_code(cls, v):
+        if v and len(v) not in {2, 5}:
+            raise ValueError("Language code must be 2 or 5 characters")
+        return v
+
+class QARequest(BaseModel):
+    question: constr(min_length=MIN_QUESTION_LENGTH, max_length=MAX_QUESTION_LENGTH)
+    language: constr(min_length=2, max_length=2) = DEFAULT_LANGUAGE
+
+    @validator('language')
+    def validate_language(cls, v):
+        if v.lower() not in SUPPORTED_LANGUAGES:
+            raise ValueError(f"Unsupported language. Supported: {SUPPORTED_LANGUAGES}")
+        return v.lower()
+
+class ErrorResponse(BaseModel):
+    error: str
+    success: bool = False
+    status_code: int
+    timestamp: str
+    details: Optional[dict] = None
+
+# Exception Handler
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request, exc):
+    error_response = ErrorResponse(
+        error=exc.detail,
+        status_code=exc.status_code,
+        timestamp=datetime.now().isoformat(),
+        details=getattr(exc, 'details', None)
+    )
+    return JSONResponse(
+        status_code=exc.status_code,
+        content=jsonable_encoder(error_response)
+    )
+
+# Helper Functions
+def get_model(model_name: str, task: str) -> Pipeline:
+    """Get or load a Hugging Face model with caching."""
+    cache_key = f"{model_name}_{task}"
+    if cache_key not in models_cache:
+        try:
+            logger.info(f"Loading model: {model_name} for task: {task}")
+            models_cache[cache_key] = pipeline(task, model=model_name)
+        except Exception as e:
+            logger.error(f"Model loading failed: {str(e)}")
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                detail="Model service unavailable",
+                details={"model": model_name, "error": str(e)}
+            )
+    return models_cache[cache_key]
+
+async def validate_and_read_file(file: UploadFile) -> Tuple[str, bytes]:
+    """Validate and read uploaded file."""
+    # Check file extension
+    file_ext = Path(file.filename).suffix[1:].lower()
+    if file_ext not in SUPPORTED_FILE_TYPES:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Unsupported file type. Supported: {', '.join(SUPPORTED_FILE_TYPES.values())}"
+        )
+
+    # Read and check file size
+    content = await file.read()
+    if len(content) > MAX_FILE_SIZE:
+        raise HTTPException(
+            status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
+            detail=f"File exceeds maximum size of {MAX_FILE_SIZE//1024//1024}MB"
+        )
 
+    await file.seek(0)
+    return file_ext, content
 
-    
-def extract_text_from_file(file_content: bytes, file_ext: str):
-    text = ""
+def extract_text(content: bytes, file_ext: str) -> str:
+    """Extract text from various file formats."""
     try:
         if file_ext == "docx":
-            doc = Document(io.BytesIO(file_content))
-            text = " ".join([p.text for p in doc.paragraphs if p.text.strip()])
-        elif file_ext in ["xls", "xlsx"]:
-            df = pd.read_excel(io.BytesIO(file_content))
-            text = " ".join(df.iloc[:, 0].dropna().astype(str).tolist())
+            doc = Document(io.BytesIO(content))
+            return " ".join(p.text for p in doc.paragraphs if p.text.strip())
+        
+        elif file_ext in {"xls", "xlsx"}:
+            df = pd.read_excel(io.BytesIO(content))
+            return " ".join(df.iloc[:, 0].dropna().astype(str).tolist())
+        
         elif file_ext == "pptx":
-            ppt = Presentation(io.BytesIO(file_content))
-            text = " ".join([shape.text for slide in ppt.slides for shape in slide.shapes if hasattr(shape, "text")])
+            ppt = Presentation(io.BytesIO(content))
+            return " ".join(shape.text for slide in ppt.slides 
+                          for shape in slide.shapes if hasattr(shape, "text"))
+        
         elif file_ext == "pdf":
-            pdf = fitz.open(stream=file_content, filetype="pdf")
-            text = " ".join([page.get_text("text") for page in pdf])
-        elif file_ext in ["jpg", "jpeg", "png"]:
-            image = Image.open(io.BytesIO(file_content))
-            text = pytesseract.image_to_string(image, config='--psm 6')
-        else:
-            raise HTTPException(status_code=400, detail="Unsupported file format.")
+            pdf = fitz.open(stream=content, filetype="pdf")
+            return " ".join(page.get_text("text") for page in pdf)
+        
+        elif file_ext in {"jpg", "jpeg", "png"}:
+            image = Image.open(io.BytesIO(content))
+            return pytesseract.image_to_string(image, config='--psm 6')
+        
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error extracting text: {str(e)}")
+        logger.error(f"Text extraction failed: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="Failed to extract text from file",
+            details={"error": str(e), "file_type": file_ext}
+        )
+
+def preprocess_text(text: str) -> str:
+    """Clean and normalize extracted text."""
+    text = re.sub(r'\s+', ' ', text).strip()
+    return text[:MAX_TEXT_LENGTH] if len(text) > MAX_TEXT_LENGTH else text
 
-    if not text.strip():
-        raise HTTPException(status_code=400, detail="No extractable text found.")
-    return text
+# API Endpoints
+@app.post("/qa")
+async def question_answering(
+    file: UploadFile = File(...),
+    question: str = Form(...),
+    language: str = Form(DEFAULT_LANGUAGE)
+) -> JSONResponse:
+    try:
+        # Validation et extraction du texte
+        file_ext, content = await validate_and_read_file(file)
+        text = preprocess_text(extract_text(content, file_ext))
+        
+        # Détection spéciale pour les questions sur le thème
+        theme_keywords = {
+            "fr": ["thème", "sujet principal", "quoi le sujet"],
+            "en": ["theme", "main topic", "what is about"]
+        }
+        
+        is_theme_question = any(
+            kw in question.lower() 
+            for kw in theme_keywords.get(language, theme_keywords["en"])
+        )
+        
+        if is_theme_question:
+            # Utilisation d'un prompt spécialisé pour l'analyse thématique
+            theme_prompt = (
+                "Extrayez le thème principal de ce texte en 1-2 phrases. "
+                "Répondez comme si vous expliquiez à un novice. "
+                "Texte : {text}"
+            )
+            
+            # Utilisation d'un LLM plus puissant pour l'analyse thématique
+            generator = get_model("moussaKam/barthez-orangesum-abstract", "text-generation")
+            response = generator(
+                theme_prompt.format(text=text[:2000]),
+                max_length=200,
+                num_return_sequences=1,
+                do_sample=False
+            )
+            
+            # Nettoyage de la réponse
+            theme = response[0]["generated_text"].split(":")[-1].strip()
+            theme = re.sub(r"^(Le|La)\s+", "", theme)  # Retire les articles en début de phrase
+            
+            return JSONResponse({
+                "question": question,
+                "answer": f"Le document traite principalement de : {theme}",
+                "confidence": 0.95,  # Haut confiance car méthode spécialisée
+                "language": language,
+                "processing_method": "theme_analysis",
+                "success": True
+            })
+        
+        # ... reste du code pour les questions normales ...
+        
+        # Standard QA processing
+        result = qa_model(question=request.question, context=clean_text)
+        
+        if result["score"] < 0.1:  # Low confidence threshold
+            return JSONResponse({
+                "question": request.question,
+                "answer": "No clear answer found in the document" if language == "en" else "Aucune réponse claire trouvée dans le document",
+                "confidence": result["score"],
+                "language": language,
+                "warning": "low_confidence",
+                "success": True
+            })
+        
+        return JSONResponse({
+            "question": request.question,
+            "answer": result["answer"],
+            "confidence": result["score"],
+            "language": language,
+            "success": True
+        })
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"QA processing failed: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Document analysis failed",
+            details={"error": str(e)}
+        )
 
     ########################################################
 @app.get("/", response_class=HTMLResponse)
@@ -160,20 +378,6 @@ async def summarize_document(file: UploadFile = File(...)):
     except Exception as e:
         raise HTTPException(500, f"Error processing document: {str(e)}")
 #################################################################
-@app.post("/qa")
-async def question_answering(file: UploadFile = File(...), question: str = Form(...)):
-    content = await file.read()
-    file_ext = file.filename.split(".")[-1].lower()
-    extracted_text = extract_text_from_file(content, file_ext)
-    # Use a pipeline as a high-level helper
-    summarizer = get_model("google-bert/bert-large-uncased-whole-word-masking-finetuned-squad", "summarization")
-    if len(extracted_text) > 2000:
-        extracted_text = summarizer(extracted_text[:2000], max_length=500, min_length=100, do_sample=False)[0]["summary_text"]
-    
-    qa_model = get_model("distilbert-base-cased-distilled-squad", "question-answering")
-    answer = qa_model(question=question, context=extracted_text)
-    
-    return {"question": question, "answer": answer["answer"], "context_used": extracted_text}
 
 ###############################################