1B_main.py

import os
import json 
import pymupdf
import numpy as np
from llama_cpp import Llama
from sentence_transformers import SentenceTransformer, util
from tqdm import tqdm
import time
import re
from collections import defaultdict
from operator import itemgetter

MODEL_PATH = "./models/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
INPUT_DIR = "./input_pdfs"
OUTPUT_DIR = "./output_json"
TOP_N_RANKED_SECTIONS = 6
TOP_N_REFINED_SECTIONS = 6

slm = Llama(model_path=MODEL_PATH, n_ctx=2048, verbose=False)
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

print("Models loaded.")

def extract_structure_from_pdf(pdf_path):
    doc = pymupdf.open(pdf_path)
    all_elements = []
    
    for page_num, page in enumerate(doc):
        blocks = page.get_text("dict", sort=True)["blocks"]
        
        page_elements = []
        for block in blocks:
            if "lines" not in block:
                continue
                
            for line in block["lines"]:
                for span in line["spans"]:
                    text = span["text"].strip()
                    if not text or len(text) < 2:
                        continue
                    
                    element = {
                        "text": text,
                        "page": page_num + 1,
                        "font_size": span["size"],
                        "is_bold": bool(span["flags"] & 2**4),
                        "bbox": span["bbox"],
                        "page_width": page.rect.width,
                        "page_height": page.rect.height
                    }
                    page_elements.append(element)
        
        sorted_page_elements = sort_reading_order(page_elements)
        all_elements.extend(sorted_page_elements)
    
    doc.close()
    return classify_complete_document_structure(all_elements, pdf_path)

def sort_reading_order(elements):
    if not elements:
        return elements
    
    lines = group_by_lines(elements, y_tolerance=10)
    
    sorted_elements = []
    for line in lines:
        sorted_line = sorted(line, key=lambda e: e["bbox"][0])
        sorted_elements.extend(sorted_line)
    
    return sorted_elements

def group_by_lines(elements, y_tolerance=10):
    if not elements:
        return []
    
    elements = sorted(elements, key=lambda e: e["bbox"][1])
    
    lines = []
    current_line = []
    last_y = None
    
    for element in elements:
        y_top = element["bbox"][1]
        
        if last_y is None or abs(y_top - last_y) <= y_tolerance:
            current_line.append(element)
            last_y = y_top if last_y is None else (last_y + y_top) / 2
        else:
            if current_line:
                lines.append(current_line)
            current_line = [element]
            last_y = y_top
    
    if current_line:
        lines.append(current_line)
    
    return lines

def classify_complete_document_structure(elements, pdf_path):
    if not elements:
        return []
    
    doc_stats = get_document_stats(elements)
    classified = []
    title_found = False
    h1_count = 0
    
    for i, element in enumerate(elements):
        classification = classify_element_complete(element, elements, i, doc_stats, title_found, h1_count)
        
        if classification["level"] == "TITLE":
            title_found = True
        elif classification["level"] == "H1":
            h1_count += 1
        
        classified.append({
            "doc_path": pdf_path,
            "page": element["page"],
            "level": classification["level"],
            "title": element["text"],
            "content": element["text"],
            "font_size": element["font_size"],
            "is_bold": element["is_bold"],
            "confidence": classification["confidence"]
        })
    
    return classified

def classify_element_complete(element, all_elements, index, doc_stats, title_found, h1_count):
    text = element["text"]
    
    if not title_found and is_document_title(element, all_elements, index, doc_stats):
        return {"level": "TITLE", "confidence": 0.9}
    
    heading_level = classify_heading(element, all_elements, index, h1_count)
    if heading_level:
        confidence = calculate_heading_confidence(element, heading_level, all_elements, index)
        return {"level": heading_level, "confidence": confidence}
    
    return {"level": "BODY", "confidence": 0.8}

def is_document_title(element, all_elements, index, doc_stats):
    text = element["text"]
    
    if element["page"] != 1:
        return False
    
    word_count = len(text.split())
    if word_count < 2 or word_count > 12:
        return False
    
    score = 0
    
    top_position = element["bbox"][1] / element["page_height"]
    if top_position < 0.15:
        score += 6
    elif top_position < 0.3:
        score += 4
    elif top_position < 0.5:
        score += 2
    
    text_width = element["bbox"][2] - element["bbox"][0]
    page_center = element["page_width"] / 2
    text_center = element["bbox"][0] + text_width / 2
    center_deviation = abs(text_center - page_center) / element["page_width"]
    
    if center_deviation < 0.1:
        score += 4
    elif center_deviation < 0.2:
        score += 2
    
    if element["is_bold"]:
        score += 3
    
    if 3 <= word_count <= 8:
        score += 3
    elif word_count <= 2:
        score += 1
    
    if is_isolated_element(element, all_elements, index):
        score += 2
    
    if re.match(r'^\d+\.', text) or text.endswith(':'):
        score -= 3
    
    title_words = ["report", "analysis", "study", "guide", "manual", "document", "overview"]
    if any(word in text.lower() for word in title_words):
        score += 2
    
    return score >= 8

def classify_heading(element, all_elements, index, h1_count):
    text = element["text"]
    word_count = len(text.split())
    
    if word_count > 15:
        return None
    
    h1_score = 0
    h2_score = 0
    h3_score = 0
    h4_score = 0
    
    if re.match(r'^\d+\.\s+[A-Z]', text):
        h1_score += 6
    elif re.match(r'^\d+\.\d+\s+[A-Z]', text):
        h2_score += 6
    elif re.match(r'^\d+\.\d+\.\d+', text):
        h3_score += 6
    elif re.match(r'^[A-Z]\.\s+[A-Z]', text):
        h2_score += 4
    elif re.match(r'^\d+\)\s+[A-Z]', text):
        h2_score += 4
    
    major_sections = ["introduction", "background", "methodology", "results", "conclusion", "references", "abstract", "summary"]
    if any(text.lower().strip() == section or text.lower().startswith(section + " ") for section in major_sections):
        h1_score += 5
    
    subsection_keywords = ["overview", "definition", "scope", "design", "evaluation", "analysis", "discussion"]
    if any(keyword in text.lower() for keyword in subsection_keywords):
        h2_score += 3
    
    if element["is_bold"]:
        h1_score += 2
        h2_score += 3
        h3_score += 2
    
    if word_count <= 3:
        h1_score += 2
        h2_score += 1
    elif word_count <= 6:
        h1_score += 1
        h2_score += 3
        h3_score += 2
    elif word_count <= 8:
        h2_score += 1
        h3_score += 3
        h4_score += 2
    elif word_count <= 10:
        h3_score += 1
        h4_score += 2
    
    if is_isolated_element(element, all_elements, index):
        h1_score += 2
        h2_score += 3
        h3_score += 2
        h4_score += 1
    
    if h1_count >= 5:
        h1_score -= 2
    
    if text.endswith('?') and word_count <= 8:
        h2_score += 2
        h3_score += 3
    
    if text.endswith(':') and word_count <= 6:
        h3_score += 3
        h4_score += 2
    
    if text.isupper() and len(text) > 3:
        h1_score += 1
        h2_score += 2
        h3_score += 1
    
    scores = {"H1": h1_score, "H2": h2_score, "H3": h3_score, "H4": h4_score}
    max_score = max(scores.values())
    
    if max_score < 4:
        return None
    
    return max(scores.keys(), key=lambda k: scores[k])

def is_isolated_element(element, all_elements, index):
    if index == 0 or index == len(all_elements) - 1:
        return True
    
    current = element
    isolation_threshold = 12
    
    if index > 0:
        prev = all_elements[index - 1]
        if prev["page"] == current["page"]:
            vertical_gap = current["bbox"][1] - prev["bbox"][3]
            if vertical_gap > isolation_threshold:
                return True
    
    if index < len(all_elements) - 1:
        next_elem = all_elements[index + 1]
        if next_elem["page"] == current["page"]:
            vertical_gap = next_elem["bbox"][1] - current["bbox"][3]
            if vertical_gap > isolation_threshold:
                return True
    
    return False

def calculate_heading_confidence(element, level, all_elements, index):
    base_confidence = 0.7
    
    if re.match(r'^\d+\.\s+[A-Z]', element["text"]):
        base_confidence += 0.2
    
    if is_isolated_element(element, all_elements, index):
        base_confidence += 0.1
    
    if element["is_bold"]:
        base_confidence += 0.1
    
    return min(base_confidence, 1.0)

def get_document_stats(elements):
    return {
        "total_elements": len(elements),
        "bold_count": sum(1 for e in elements if e["is_bold"]),
        "avg_length": sum(len(e["text"]) for e in elements) / len(elements) if elements else 0
    }

def process_all_pdfs_in_folder():
    pdf_files = []
    
    for filename in sorted(os.listdir(INPUT_DIR)):
        if filename.lower().endswith('.pdf'):
            pdf_path = os.path.join(INPUT_DIR, filename)
            pdf_files.append(pdf_path)
    
    if not pdf_files:
        print(f"No PDF files found in '{INPUT_DIR}'.")
        return
    
    print(f"Found {len(pdf_files)} PDF files to process.")
    
    start_time = time.time()
    all_results = []
    
    for pdf_path in pdf_files:
        print(f"Processing: {os.path.basename(pdf_path)}")
        extracted_sections = extract_structure_from_pdf(pdf_path)
        
        for section in extracted_sections:
            all_results.append({
                "document": os.path.basename(section["doc_path"]),
                "page": section["page"],
                "level": section["level"],
                "text": section["title"],
                "font_size": section["font_size"],
                "is_bold": section["is_bold"],
                "confidence": section["confidence"]
            })
    
    pdf_order = [os.path.basename(path) for path in pdf_files]
    pdf_pos_map = {pdf_name: i for i, pdf_name in enumerate(pdf_order)}
    
    sorted_results = sorted(
        all_results,
        key=lambda x: (pdf_pos_map.get(x['document'], float('inf')), x['page'])
    )
    
    output_data = {
        "metadata": {
            "total_pdfs_processed": len(pdf_files),
            "total_elements_classified": len(sorted_results),
            "processing_time_seconds": round(time.time() - start_time, 2),
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        },
        "extracted_structure": sorted_results
    }
    
    output_filename = f"complete_structure_{time.strftime('%Y%m%d_%H%M%S')}.json"
    output_path = os.path.join(OUTPUT_DIR, output_filename)
    
    with open(output_path, 'w') as f:
        json.dump(output_data, f, indent=2)
    
    level_counts = defaultdict(int)
    for result in sorted_results:
        level_counts[result["level"]] += 1
    
    print(f"\nPDFs processed: {len(pdf_files)}")
    print(f"Total elements classified: {len(sorted_results)}")
    print(f"Processing time: {output_data['metadata']['processing_time_seconds']} seconds")
    print(f"Results saved to: {output_path}")
    
    for level in ["TITLE", "H1", "H2", "H3", "H4", "BODY"]:
        if level in level_counts:
            print(f"{level}: {level_counts[level]} elements")
    
    for i, result in enumerate(sorted_results[:20]):
        if i == 0 or result["document"] != sorted_results[i-1]["document"]:
            print(f"\n{result['document']}")
        
        print(f"[{result['level']}] {result['text'][:60]}... (p{result['page']})")

def ranking(sections, persona, jtbd, embedding_model_instance):
    if not sections:
        return []
    
    query = f"As a {persona}, I need to {jtbd}."
    query_embedding = embedding_model_instance.encode(query, convert_to_tensor=True)
    section_contents = [section['content'] for section in sections]
    section_emb = embedding_model_instance.encode(section_contents, convert_to_tensor=True)
    similarity = util.cos_sim(query_embedding, section_emb)[0].cpu().numpy()

    for i, section in enumerate(sections):
        section['similarity'] = float(similarity[i])

    return sorted(sections, key=lambda x: x['similarity'], reverse=True)

def refine_sections(ranked_sections, persona, jtbd, slm_instance, n):
    if not ranked_sections:
        return []
    
    prompt_template = """You are a {persona}. Your goal is to {jtbd}.

Based on this context, provide a concise, one-sentence summary of the key insight from the following text that would be most relevant for your goal.

Text: {section_text}

Key insight:"""
    
    result = []
    for i, section in enumerate(tqdm(ranked_sections[:n], desc="Refining Sections", leave=False)):
        try:
            section_text = section.get("content", section.get("title", "")).strip()
            if not section_text:
                print(f"Warning: Empty content for section {i+1}")
                result.append({
                    "document": os.path.basename(section["doc_path"]),
                    "page_number": section["page"],
                    "refined_text": "No content available for analysis."
                })
                continue
            
            prompt = prompt_template.format(
                persona=persona, 
                jtbd=jtbd, 
                section_text=section_text[:1000]
            )
            
            output = slm_instance(
                prompt, 
                max_tokens=200,
                stop=["\n\n", "Key insight:", "Text:"],
                echo=False,
                temperature=0.7,
                top_p=0.9
            )
            
            refined_text = output['choices'][0]['text'].strip()
            
            if not refined_text:
                output_fallback = slm_instance(
                    prompt, 
                    max_tokens=100,
                    echo=False,
                    temperature=0.5
                )
                refined_text = output_fallback['choices'][0]['text'].strip()
            
            refined_text = clean_response(refined_text)
            
            if not refined_text:
                refined_text = f"This section contains information about {section.get('title', 'the topic')} that may be relevant for trip planning."
            
            result.append({
                "document": os.path.basename(section["doc_path"]),
                "page_number": section["page"],
                "refined_text": refined_text
            })
            
        except Exception as e:
            print(f"Error processing section {i+1}: {str(e)}")
            result.append({
                "document": os.path.basename(section["doc_path"]),
                "page_number": section["page"],
                "refined_text": f"Error processing this section: {str(e)}"
            })
    
    return result

def clean_response(text):
    if not text:
        return ""
    
    text = text.replace("Key insight:", "").strip()
    text = text.replace("Summary:", "").strip()
    
    text = re.sub(r'^[-•*]\s*', '', text)
    
    if text and not text.endswith(('.', '!', '?')):
        text += "."
    
    if text:
        text = text[0].upper() + text[1:]
    
    return text

def main():
    input_json_path = ""
    for f in os.listdir(INPUT_DIR):
        if f.lower().endswith(".json"):
            input_json_path = os.path.join(INPUT_DIR, f)
            break
    
    if not input_json_path:
        print(f"No input JSON file found in '{INPUT_DIR}'.")
        return

    with open(input_json_path, 'r') as f:
        input_data = json.load(f)

    persona = input_data.get("persona", {}).get("role", "user")
    job_to_be_done = input_data.get("job_to_be_done", {}).get("task", "analyze documents")
    documents = input_data.get("documents", [])
    pdf_files = [os.path.join(INPUT_DIR, doc["filename"]) for doc in documents]

    start_time = time.time()

    all_extracted_sections = []
    for pdf_path in pdf_files:
        all_extracted_sections.extend(extract_structure_from_pdf(pdf_path))

    ranked_sections = ranking(all_extracted_sections, persona, job_to_be_done, embedding_model)

    extracted_sections_output = []
    for i, section in enumerate(ranked_sections[:TOP_N_RANKED_SECTIONS]):
        extracted_sections_output.append({
            "document": os.path.basename(section["doc_path"]),
            "section_title": section["title"],
            "importance_rank": i + 1,
            "page_number": section["page"]
        })

    subsection_analysis = refine_sections(ranked_sections, persona, job_to_be_done, slm, TOP_N_REFINED_SECTIONS)

    output_data = {
        "metadata": {
            "input_documents": [doc["filename"] for doc in documents],
            "persona": persona,
            "job_to_be_done": job_to_be_done,
            "processing_timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        },
        "extracted_sections": extracted_sections_output,
        "subsection_analysis": subsection_analysis
    }

    output_filename = f"{input_data['challenge_info']['test_case_name']}.json"
    output_path = os.path.join(OUTPUT_DIR, output_filename)
    
    with open(output_path, 'w') as f:
        json.dump(output_data, f, indent=4)

    total_time = time.time() - start_time
    print(f"Total processing time: {total_time:.2f} seconds.")
    print(f"Results saved to: {output_path}")

if __name__ == "__main__":
    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)
    
    main()