medical-agent / template_analyser_llm.py
Nourhenem's picture
Upload folder using huggingface_hub
1eb76aa verified
#!/usr/bin/env python3
# coding: utf-8
"""
Test file: Section extraction via GPT for medical Word templates
This script reads a .docx document, sends the entire document to GPT for classification
into medical sections, and extracts the structured result.
"""
import os
import json
from docx import Document
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
# 1. Configure your OpenAI API key in the environment
# export OPENAI_API_KEY="sk-..."
api_key = os.getenv('OPENAI_API_KEY')
# 2. Define the prompt for section classification
section_prompt = ChatPromptTemplate.from_messages([
("system", """
Vous êtes un expert en analyse de documents médicaux. Je vais vous fournir le texte complet d'un rapport médical.
Votre tâche est de :
1. Identifier automatiquement toutes les sections dans le document. Une section est définie par :
- Un en-tête qui peut être en majuscules ou en titre (ex: "SCANNER", "IRM DU GENOU DROIT")
- Un titre de section suivi de deux points ":" (ex: "Indication:", "Technique:", "Résultats:", "Conclusions:", etc.)
- Les sections peuvent aussi être "Indications", "Techniques", "CONCLUSION" (sans deux points)
2. Pour chaque section identifiée, extraire son contenu en collectant toutes les lignes suivantes jusqu'à la prochaine section.
3. Identifier les champs à remplir par l'utilisateur :
- Les balises <ASR_VOX> indiquent des champs à remplir
- Les textes génériques comme "xxx", "xxxx" indiquent des champs à remplir
- Les formules conditionnelles comme "SI(Civilité Nom usuel médecin..." indiquent des champs à remplir
4. Retourner un objet JSON valide avec cette structure exacte :
{{
"document_type": "type de document détecté",
"sections": {{
"nom_section": {{
"content": "contenu brut de la section",
"has_user_fields": true,
"user_fields": ["liste des champs à remplir"]
}}
}}
}}
Répondez UNIQUEMENT avec le JSON—aucun commentaire supplémentaire.
"""),
("human", "Voici le texte complet du rapport médical :\n\n{document_text}\n\nExtrayez toutes les sections et identifiez les champs à remplir.")
])
# 3. Initialize the LLM with appropriate parameters
llm = ChatOpenAI(
model="gpt-4o-mini",
temperature=0,
max_tokens=4000
)
# 4. Create the chain using new LangChain API
section_classifier = section_prompt | llm
def extract_sections_via_gpt(docx_path: str):
"""Extracts and returns structured sections from the entire document via GPT."""
if not os.path.exists(docx_path):
raise FileNotFoundError(f"Document not found: {docx_path}")
# Read the entire document
doc = Document(docx_path)
# Combine all paragraphs into one text
document_text = ""
for para in doc.paragraphs:
text = para.text.strip()
if text:
document_text += text + "\n"
if not document_text.strip():
return {"error": "Document appears to be empty"}
print(f"Document text preview: {document_text[:200]}...")
try:
# Send entire document to GPT for processing using new API
response = section_classifier.invoke({"document_text": document_text})
# Extract content from the response
result = response.content.strip()
print(f"GPT response: {result[:500]}...")
# Parse the JSON response
sections_data = json.loads(result)
return sections_data
except json.JSONDecodeError as e:
print(f"Erreur de parsing JSON: {e}")
print(f"Réponse brute: {result}")
return {"error": f"Invalid JSON response: {e}", "raw_response": result}
except Exception as e:
print(f"Erreur lors de l'appel GPT: {e}")
return {"error": f"GPT processing error: {e}"}
def print_results(sections_data):
"""Print the extracted sections in a readable format."""
if "error" in sections_data:
print(f"Erreur: {sections_data['error']}")
if "raw_response" in sections_data:
print(f"Réponse brute: {sections_data['raw_response']}")
return
print("=" * 50)
print("ANALYSE DU DOCUMENT MÉDICAL")
print("=" * 50)
if "document_type" in sections_data:
print(f"Type de document: {sections_data['document_type']}")
print()
if "sections" in sections_data:
for section_name, section_data in sections_data["sections"].items():
print(f"📋 SECTION: {section_name}")
print("-" * 30)
if isinstance(section_data, dict):
if "content" in section_data:
print(f"Contenu: {section_data['content']}")
if section_data.get("has_user_fields", False):
print("⚠️ Cette section contient des champs à remplir par l'utilisateur:")
for field in section_data.get("user_fields", []):
print(f" • {field}")
else:
print("✅ Cette section est complète")
else:
print(f"Contenu: {section_data}")
print()
if __name__ == '__main__':
# Path to your sample .docx
SAMPLE_DOCX = 'sample.docx'
try:
print("Extraction des sections en cours...")
sections = extract_sections_via_gpt(SAMPLE_DOCX)
# Print formatted results
print_results(sections)
# Also save raw JSON for debugging
with open('extracted_sections.json', 'w', encoding='utf-8') as f:
json.dump(sections, f, indent=2, ensure_ascii=False)
print("\nRésultats sauvegardés dans 'extracted_sections.json'")
except Exception as e:
print(f"Erreur principale: {e}")
import traceback
traceback.print_exc()