smart-models/Placebo_AI

1import os2import json3import sys4import fitz5import requests6import base647from tqdm import tqdm8from PIL import Image9 10# Configure Tesseract path if needed11import pytesseract12pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'13os.environ["PATH"] += os.pathsep + r'C:\Program Files\Tesseract-OCR'14 15OLLAMA_URL = "http://localhost:11434/api/generate"16 17def check_ollama_status():18    """Checks if Ollama is running and has moondream loaded."""19    try:20        response = requests.get("http://localhost:11434/api/tags", timeout=5)21        if response.status_code == 200:22            models = [m.get("name") for m in response.json().get("models", [])]23            print(f"Ollama is running. Available models: {models}")24            return "moondream:latest" in models or "moondream" in models25    except Exception as e:26        print(f"Warning: Could not connect to Ollama at localhost:11434. VLM descriptions will be skipped. ({e})")27    return False28 29def get_vlm_description(image_path):30    """Call Ollama VLM (Moondream) for image description."""31    prompt = (32        "Describe any anatomical diagrams, histology slides, physiological pathways, clinical charts, "33        "flowcharts, medical algorithms, or dosage tables on this page. If you see a process or cycle, "34        "extract it as a step-by-step logical sequence (Step 1 -> Step 2). Keep descriptions concise and medical-grade. "35        "Do not write long paragraphs. If none, say 'No visual data'."36    )37    try:38        with open(image_path, "rb") as image_file:39            img_str = base64.b64encode(image_file.read()).decode('utf-8')40        41        payload = {42            "model": "moondream",43            "prompt": prompt,44            "images": [img_str],45            "stream": False46        }47        48        response = requests.post(OLLAMA_URL, json=payload, timeout=90)49        if response.status_code == 200:50            return response.json().get("response", "").strip()51        else:52            return f"Ollama VLM Error: {response.status_code}"53    except Exception as e:54        return f"VLM skipped: {e}"55 56def clean_filename(name):57    return name.replace(" ", "_").replace("&", "and").replace("'", "").replace("(", "").replace(")", "").replace("[", "").replace("]", "").replace(",", "")58 59def process_books(pilot=False):60    mbbs_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "raw", "MBBS")61    master_jsonl_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "processed")62    quality_report_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config", "mbbs_quality_report.json")63    64    # Check Ollama65    has_vlm = check_ollama_status()66    67    # Load quality report68    quality_report = {}69    if os.path.exists(quality_report_path):70        print(f"Loading quality report from {quality_report_path}...")71        try:72            with open(quality_report_path, "r", encoding="utf-8") as f:73                quality_report = json.load(f)74            print(f"Loaded quality profiles for {len(quality_report)} books.")75        except Exception as e:76            print(f"Error reading quality report: {e}. Will fall back to dynamic page detection.")77    else:78        print("Warning: Quality report not found! Will fall back to dynamic page checks.")79 80    # Gather PDFs81    pdf_files = []82    83    allowed_subjects = [84        "ENT",85        "Embryology",86        "NEUROSCIENCE",87    ]88    89    for root, dirs, files in os.walk(mbbs_dir):90        for file in files:91            if file.lower().endswith(".pdf"):92                full_path = os.path.join(root, file)93                rel = os.path.relpath(root, mbbs_dir)94                parts = rel.split(os.sep)95                if len(parts) >= 2 and parts[0] == "books":96                    subject = parts[1]97                else:98                    subject = rel if rel != "." else "general"99                100                if subject not in allowed_subjects:101                    continue102                103                pdf_files.append({104                    "subject": subject,105                    "book": file,106                    "path": full_path107                })108                109    pdf_files = sorted(pdf_files, key=lambda x: (x["subject"], x["book"]))110    111    if pilot:112        print("\n=== PILOT MODE ACTIVE ===")113        print("Restricting processing to 1 book per subject, and 2 clean pages per book.")114        # Group by subject and pick first book115        pilot_books = {}116        for b in pdf_files:117            if b["subject"] not in pilot_books:118                pilot_books[b["subject"]] = b119        pdf_files = list(pilot_books.values())120        print(f"Pilot set contains {len(pdf_files)} books.")121 122    # Process books subject-wise123    # To keep files separate, we open output file descriptors for each subject124    subject_files = {}125    126    for idx, book_info in enumerate(pdf_files):127        book_name = book_info["book"]128        subject = book_info["subject"]129        path = book_info["path"]130        131        book_clean = clean_filename(book_name.replace(".pdf", ""))132        subject_clean = clean_filename(subject)133        134        # Setup output directories for page images135        # Images MUST be saved inside 'd:\sample chatbot\data' for fastapi to serve them136        img_out_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "mbbs_processed", subject_clean, book_clean)137        os.makedirs(img_out_dir, exist_ok=True)138        139        # Open output JSONL file140        jsonl_name = f"mbbs_{subject_clean}.jsonl"141        jsonl_path = os.path.join(master_jsonl_dir, jsonl_name)142        143        # If running full, we append. In pilot, we overwrite.144        mode = "w" if pilot else "a"145        146        print(f"\n[{idx+1}/{len(pdf_files)}] Processing book: {book_name} under subject: {subject}")147        148        # Retrieve excluded pages149        book_profile = quality_report.get(book_name, {})150        excluded_pages = set()151        if book_profile:152            excluded_pages.update(book_profile.get("blank_pages", []))153            excluded_pages.update(book_profile.get("toc_pages", []))154            excluded_pages.update(book_profile.get("index_pages", []))155            excluded_pages.update(book_profile.get("copyright_preface_pages", []))156            excluded_pages.update(book_profile.get("blurry_pages", []))157            print(f"  - Quality report: Excluded {len(excluded_pages)} non-useful pages (blurry, blank, indices, TOC, copyright).")158        159        try:160            doc = fitz.open(path)161            total_pages = doc.page_count162        except Exception as e:163            print(f"  Error opening PDF {book_name}: {e}")164            continue165            166        processed_count = 0167        168        # Write to JSONL169        with open(jsonl_path, mode, encoding="utf-8") as f_out:170            for page_idx in tqdm(range(total_pages), desc=f"Processing {book_clean[:30]}"):171                page_num = page_idx + 1172                173                # Check exclusions174                if page_num in excluded_pages:175                    continue176                177                # Dynamic check fallback if quality report was missing178                page = doc[page_idx]179                text = page.get_text().strip()180                text_len = len(text)181                182                if not book_profile:183                    # Basic dynamic filters184                    if text_len < 50 and len(page.get_images()) == 0:185                        continue # Blank186                    # Simple keyword checks for TOC / Index187                    text_lower = text.lower()188                    if "table of contents" in text_lower or ("index" in text_lower and page_num > total_pages - 10):189                        continue # Skip index/TOC190                191                # Extract clean page content192                # If page is scanned, we can run OCR (Pytesseract), otherwise PyMuPDF text is used193                combined_content = text194                has_images = len(page.get_images()) > 0195                196                # Render page image for the UI and VLM197                img_name = f"page_{page_num:04d}.png"198                img_path = os.path.join(img_out_dir, img_name)199                200                # Render page to file201                try:202                    pix = page.get_pixmap(dpi=150) # high enough quality for UI reading203                    pix.save(img_path)204                except Exception as ex:205                    print(f"  Warning: failed to render page image {page_num}: {ex}")206                207                # Run OCR if there's no text but there is a scanned image208                if text_len < 100 and has_images:209                    try:210                        ocr_text = pytesseract.image_to_string(Image.open(img_path))211                        if ocr_text.strip():212                            combined_content = ocr_text213                    except Exception as ocr_err:214                        combined_content = f"OCR Failed: {ocr_err}"215                216                # Add VLM description if page has images and Ollama is active217                vlm_desc = "No visual data"218                if has_images and has_vlm:219                    vlm_desc = get_vlm_description(img_path)220                221                if vlm_desc and vlm_desc != "No visual data":222                    combined_content += f"\n\nVISUAL DATA:\n{vlm_desc}"223                224                # Construct page dictionary225                page_data = {226                    "subject": subject,227                    "book_name": book_name,228                    "page_number": f"{page_num:04d}",229                    "image_path": os.path.abspath(img_path),230                    "content": combined_content,231                    "category": subject232                }233                234                f_out.write(json.dumps(page_data) + "\n")235                f_out.flush()236                237                processed_count += 1238                239                if pilot and processed_count >= 2:240                    break241        242        doc.close()243        print(f"  Successfully processed {processed_count} pages for {book_name}.")244 245    print("\nProcessing session completed.")246 247if __name__ == "__main__":248    is_pilot = "--pilot" in sys.argv or True # default to pilot mode unless specified, wait let's make it check argv249    # Let's inspect argv to determine pilot250    is_pilot = "--full" not in sys.argv251    process_books(pilot=is_pilot)252