SLM_Approach.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

MODEL_PATH = r"C:\Users\sahni\OneDrive\Documents\GitHub\Adobe-India-Hackathon-1B\models\tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
INPUT_DIR = r"C:\Users\sahni\OneDrive\Documents\GitHub\Adobe-India-Hackathon-1B\input_pdfs"
OUTPUT_DIR = r"C:\Users\sahni\OneDrive\Documents\GitHub\Adobe-India-Hackathon-1B\output_json"
TOP_N_RANKED_SECTIONS  =5
TOP_N_REFINED_SECTIONS = 5

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

print("Models loaded.")

# Part 1: Extraction using SLM 

def extract_structure_from_pdf(pdf_path, slm_instance):
    """
    Uses an SLM to extract headings (H1, H2, H3) from each page of a PDF.
    This version escapes the curly braces in the f-string prompt.
    """
    print(f"  [Stage 1] Extracting structure from {os.path.basename(pdf_path)}.")
    doc = pymupdf.open(pdf_path)
    all_headings = []

    # prompt with escaped curly braces {{...}} for the example
    structure_prompt_template = """Analyze the following document text. Identify only H1, H2, and H3 headings.
Do not invent headings. Output ONLY a valid JSON list of objects. Each object must have "level" and "text" keys.
Ensure all objects in the list are separated by a comma.
Example: [{{"level": "H1", "text": "Introduction"}}, {{"level": "H2", "text": "Background"}}]

Text:
---
{page_text}
---

JSON Output:
"""

    for page_num, page in enumerate(tqdm(doc, desc="  Processing pages", leave=False)):
        page_text = page.get_text("text")
        if not page_text.strip():
            continue

        # The .format() method correctly fills in {page_text} without being confused by the {{...}}
        prompt = structure_prompt_template.format(page_text=page_text)
        
        try:
            output = slm_instance(prompt, max_tokens=512, stop=["]"], echo=False)
            raw_json = output['choices'][0]['text'] + "]"
            
            json_start = raw_json.find('[')
            if json_start != -1:
                raw_json = raw_json[json_start:]
            
            cleaned_json = re.sub(r'}\s*{', '},{', raw_json)
            headings_on_page = json.loads(cleaned_json)
            
            for heading in headings_on_page:
                if isinstance(heading, dict) and "level" in heading and "text" in heading:
                    all_headings.append({
                        "doc_path": pdf_path,
                        "page": page_num + 1,
                        "level": heading["level"],
                        "title": heading["text"],
                        "content": heading["text"] 
                    })
        except (json.JSONDecodeError, IndexError, KeyError):
            pass
            
    return all_headings



# Part 2: Similarity Semantics and Embeddings 

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)


# Part 3: Refining the Sections using LM

def refine_sections(ranked_sections, persona, jtbd, slm_instance, n):
    if not ranked_sections:
        return []
    
    prompt_template = """ Your persona is a '{persona}'. Your goal is to '{jtbd}'. Based on this, provide a concise, one-sentence summary of the key insight from the follwing text.

Text: 

{section_text}

Summary:
"""
    result = []
    for section in tqdm(ranked_sections[:n], desc="Refining Sections", leave = False):
        prompt = prompt_template.format(persona = persona, jtbd = jtbd, section_text = section["content"])
        output = slm_instance(prompt, max_tokens = 150, stop=["\n", "."], echo = False)
        refined_text = output['choices'][0]['text'].strip()
        result.append({
            "document": os.path.basename(section["doc_path"]),
            "page_number": section["page"],
            "refined_text": refined_text
        })

    return result


# Final Function 

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, slm))

    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()