Update app.py

app.py

@@ -1,83 +1,119 @@
 import gradio as gr
-import PyPDF2
 import torch
+import os
+import pytesseract
+import cv2
+import datetime
+import shutil
+from pdf2image import convert_from_path
+from PIL import Image
+from tqdm import tqdm
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# --- CONFIGURATION & MODÈLES ---
 MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+OCR_LANG = "fra+ara"  # Support Français + Arabe
+ARCHIVE_DIR = '/content/archive'
+QR_DIR = '/content/processed_qrcodes'
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-
-# Charger modèle sans device_map
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME
-)
+os.makedirs(ARCHIVE_DIR, exist_ok=True)
+os.makedirs(QR_DIR, exist_ok=True)
 
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
-
-def extract_text(pdf_file):
-    if pdf_file is None:
-        return ""
-
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
+
+# --- POLITIQUES DE RÉTENTION ---
+RETENTION_POLICY = {
+    "ressources humaines": {"active": 5, "semi": 10, "archived": float('inf')},
+    "facture": {"active": 7, "semi": 3, "archived": float('inf')},
+    "certificat medicale": {"active": 2, "semi": 0, "archived": float('inf')},
+    "autre": {"active": 1, "semi": 0, "archived": float('inf')}
+}
+
+# --- FONCTIONS OCR & CLASSIFICATION ---
+
+def extract_text_advanced(file_path):
+    """Extraction OCR optimisée pour PDF scannés (Français/Arabe)."""
+    if not file_path: return ""
+    text = ""
     try:
-        reader = PyPDF2.PdfReader(pdf_file.name)
-        text = ""
-        for page in reader.pages:
-            text += (page.extract_text() or "") + "\n"
+        # Conversion PDF en images haute résolution pour l'OCR
+        pages = convert_from_path(file_path, dpi=400)
+        for i, page in enumerate(pages):
+            img_path = f"/tmp/page_{i}.png"
+            page.save(img_path, "PNG")
+            img = cv2.imread(img_path)
+            # Prétraitement de l'image pour améliorer la lecture
+            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+            gray = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]
+            text += pytesseract.image_to_string(gray, lang=OCR_LANG) + "\n"
+            os.remove(img_path)
         return text.strip()
     except Exception as e:
-        return f"Erreur PDF: {str(e)}"
-
-def chat(message, history, pdf_file):
+        return f"Erreur OCR: {str(e)}"
+
+def classify_doc(text):
+    """Classification automatique basée sur les mots-clés."""
+    text_lower = text.lower()
+    if "certificat médical" in text_lower or "منحة مرض" in text_lower:
+        return "certificat medicale"
+    elif "facture" in text_lower:
+        return "facture"
+    elif "ressources humaines" in text_lower or "rh" in text_lower:
+        return "ressources humaines"
+    return "autre"
+
+# --- LOGIQUE DE CHAT ---
+
+def chat_interface(message, history, pdf_file):
     history = history or []
-
-    pdf_text = extract_text(pdf_file)
-    if not pdf_text.strip():
-        pdf_text = "Aucun texte détecté dans le PDF."
-
-    context = pdf_text[:2500]
-
-    prompt = f"""
-Tu es ArchivChat.
-Réponds uniquement avec les informations du document.
-
-DOCUMENT:
-{context}
-
-QUESTION:
-{message}
-
-REPONSE:
-"""
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-
+    pdf_text = extract_text_advanced(pdf_file.name) if pdf_file else ""
+    category = classify_doc(pdf_text)
+    
+    # Instructions personnalisées (AI Covers & Radio) + Contexte document
+    system_prompt = (
+        "Tu es ArchivChat. Tu es expert en Radio et AI Covers. "
+        f"Document analysé (Catégorie: {category}). "
+        "Réponds en priorité selon le DOCUMENT. Si absent, utilise tes connaissances AI/Radio."
+    )
+    
+    context = pdf_text[:2000]
+    full_prompt = f"<|system|>\n{system_prompt}\nDOC:{context}</s>\n"
+    for h in history:
+        full_prompt += f"<|user|>\n{h[0]}</s>\n<|assistant|>\n{h[1]}</s>\n"
+    full_prompt += f"<|user|>\n{message}</s>\n<|assistant|>\n"
+
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
     with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=200,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
-
-    answer = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    if "REPONSE:" in answer:
-        answer = answer.split("REPONSE:")[-1].strip()
-
-    history.append((message, answer))
-    return "", history
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 📁 ArchivChat (Sans clé API)")
-
+        output = model.generate(**inputs, max_new_tokens=200, temperature=0.7)
+    
+    response = tokenizer.decode(output[0], skip_special_tokens=True).split("<|assistant|>")[-1].strip()
+    history.append((message, response))
+    
+    # Déplacement automatique vers l'archive après traitement
+    if pdf_file:
+        dest = os.path.join(ARCHIVE_DIR, os.path.basename(pdf_file.name))
+        shutil.copy(pdf_file.name, dest)
+        
+    return "", history, f"Catégorie détectée : {category.upper()}"
+
+# --- INTERFACE GRADIO ---
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 📁 ArchivChat & OCR Pro")
     with gr.Row():
-        pdf_file = gr.File(label="Uploader un PDF", file_types=[".pdf"])
-        chatbot = gr.Chatbot()
-
-    msg = gr.Textbox(label="Pose ta question")
-    msg.submit(chat, [msg, chatbot, pdf_file], [msg, chatbot])
-
-demo.launch()
-
+        with gr.Column(scale=1):
+            file_input = gr.File(label="Uploader PDF/Image")
+            status_label = gr.Label(label="Analyse Document")
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(height=450)
+            msg = gr.Textbox(label="Question", placeholder="Posez une question sur le document ou sur la radio/AI covers...")
+            clear = gr.Button("Effacer")
+
+    msg.submit(chat_interface, [msg, chatbot, file_input], [msg, chatbot, status_label])
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+if __name__ == "__main__":
+    # Note : Nécessite l'installation système : apt install tesseract-ocr-fra tesseract-ocr-ara poppler-utils
+    demo.launch(debug=True)