D1 02 Prompt templates and parsing

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
from langchain_huggingface.llms import HuggingFacePipeline
from langchain_huggingface import ChatHuggingFace
from langchain_core.prompts import PromptTemplate
from langchain_core.prompts.chat import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate, AIMessagePromptTemplate
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel
from typing import List
import json
import re
import os

# Download model from HuggingFace (same base model as D1_01)
HF_LLM_MODEL = "NousResearch/Nous-Hermes-2-Mistral-7B-DPO"

# 4-bit quantization config for efficient loading
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(HF_LLM_MODEL)

# Load model with 4-bit quantization
model = AutoModelForCausalLM.from_pretrained(
    HF_LLM_MODEL,
    device_map="auto",
    quantization_config=quantization_config,
)

# Create text generation pipeline
text_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    do_sample=True,
    temperature=0.7,
    return_full_text=False,
    eos_token_id=tokenizer.eos_token_id,
)

llm = HuggingFacePipeline(pipeline=text_pipeline)

prompt_template = PromptTemplate(
    input_variables=["topic"],
    template="Explain the following topic in simple terms:\n\n{topic}"
)

print(prompt_template.format(topic="What is machine learning?"))

response = llm.invoke(prompt_template.format(topic="What is machine learning?"))
print(response)

simplify_prompt = PromptTemplate(
    input_variables=["clause"],
    template="""
You are a legal assistant that simplifies complex legal clauses into plain, understandable English.

Clause:
{clause}

Simplified Explanation:
"""
)

legal_clause = (
    "The lessee shall indemnify and hold harmless the lessor from any liabilities, damages, "
    "or claims arising out of the use of the premises, except in cases of gross negligence."
)

formatted_prompt = simplify_prompt.format(clause=legal_clause)
response = llm.invoke(formatted_prompt)

print("Simplified:\n", response)

# Ollama configuration (no API key needed!)
OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://ollama:11434")

# === Model Configuration ===
HF_LLM_MODEL = "NousResearch/Nous-Hermes-2-Mistral-7B-DPO-GGUF"
OLLAMA_LLM_MODEL = f"hf.co/{HF_LLM_MODEL}:Q4_K_M"

print(f"Ollama host: {OLLAMA_HOST}")
print(f"Model: {OLLAMA_LLM_MODEL}")

from langchain_openai import ChatOpenAI

# Use Ollama as OpenAI-compatible endpoint (no API key required)
llm_ollama = ChatOpenAI(
    base_url=f"{OLLAMA_HOST}/v1",
    api_key="ollama",  # Ollama ignores this but LangChain requires it
    model=OLLAMA_LLM_MODEL,
    temperature=0.7,
    max_tokens=512
)

response = llm_ollama.invoke(formatted_prompt)

print("Simplified:\n", response)

template = "You are a helpful assistant that translates complex legal terms into plain and understandable language."
system_message_prompt = SystemMessagePromptTemplate.from_template(template)

legal_text_1 = "Notwithstanding any provision to the contrary herein, the indemnitor agrees to indemnify, defend, and hold harmless the indemnitee from and against any and all claims, liabilities, damages, or expenses (including, without limitation, reasonable attorney’s fees) arising out of or related to the indemnitor’s acts or omissions, except to the extent that such claims, liabilities, damages, or expenses result from the gross negligence or willful misconduct of the indemnitee."
example_input_1 = HumanMessagePromptTemplate.from_template(legal_text_1)

plain_text_1 = "One party agrees to cover any costs, claims, or damages that happen because of their actions, including legal fees. However, they do not have to pay if the other party was extremely careless or acted intentionally wrong."
example_output_1 = AIMessagePromptTemplate.from_template(plain_text_1)

legal_text_2 = "This agreement shall be binding upon and inure to the benefit of the parties hereto and their respective heirs, executors, administrators, successors, and assigns, and shall not be assignable by either party without the prior written consent of the other, except that either party may assign its rights and obligations hereunder in connection with a merger, consolidation, or sale of substantially all of its assets."
example_input_2 = HumanMessagePromptTemplate.from_template(legal_text_2)

plain_text_2 = "This agreement applies to both parties and their future representatives, such as heirs or business successors. Neither party can transfer their rights under this agreement to someone else unless they get written permission. However, if one party merges with another company or sells most of its assets, they can transfer their rights without permission."
example_output_2 = AIMessagePromptTemplate.from_template(plain_text_2)

human_template = "{legal_text}"
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)

chat_prompt = ChatPromptTemplate.from_messages(
    [system_message_prompt, example_input_1, example_output_1, example_input_2, example_output_2, human_message_prompt]
)

some_example_text = "Any waiver of any term or condition of this agreement shall not be deemed a continuing waiver of such term or condition, nor shall it be considered a waiver of any other term or condition hereof. No failure or delay by either party in exercising any right, power, or privilege under this agreement shall operate as a waiver thereof, nor shall any single or partial exercise preclude any other or further exercise thereof or the exercise of any other right, power, or privilege."
request = chat_prompt.format_prompt(legal_text=some_example_text).to_messages()

result = llm_ollama.invoke(request)

print(result)

template = "You are a helpful assistant that translates complex legal terms into plain and understandable language.  Respond only with a JSON object containing a single key 'translation' and its corresponding value."
system_message_prompt = SystemMessagePromptTemplate.from_template(template)

chat_prompt = ChatPromptTemplate.from_messages(
    [system_message_prompt, example_input_1, example_output_1, example_input_2, example_output_2, human_message_prompt]
)

some_example_text = "Any waiver of any term or condition of this agreement shall not be deemed a continuing waiver of such term or condition, nor shall it be considered a waiver of any other term or condition hereof. No failure or delay by either party in exercising any right, power, or privilege under this agreement shall operate as a waiver thereof, nor shall any single or partial exercise preclude any other or further exercise thereof or the exercise of any other right, power, or privilege."
request = chat_prompt.format_prompt(legal_text=some_example_text).to_messages()

result = llm_ollama.invoke(request)

result

# Define output structure
class ClauseSummary(BaseModel):
    parties: List[str]
    start_date: str
    penalty_clause: str

# Set up parser
parser = PydanticOutputParser(pydantic_object=ClauseSummary)
format_instructions = parser.get_format_instructions()

# Create prompt with correct input variables
prompt = PromptTemplate(
    input_variables=["clause", "format_instructions"],
    template="""
Extract the following fields from the contract clause below and return them in **valid JSON format ONLY**, with no extra text or explanation.

Clause:
{clause}

{format_instructions}
"""
)

# Clause to parse
clause_text = (
    "The agreement between Acme Corp and Beta LLC begins on January 1, 2025. "
    "If either party breaks the agreement, a €5,000 penalty applies."
)

# Format the full prompt
full_prompt = prompt.format(clause=clause_text, format_instructions=format_instructions)

# Run the model
response = llm_ollama.invoke(full_prompt)

# Parse the response
try:
    parsed = parser.parse(response.content)
    print(parsed.model_dump())
except Exception as e:
    print("Could not parse output.")
    print("Raw response:", response)
    print(e)

print(parsed.parties)
print(parsed.start_date)
print(parsed.penalty_clause)

# Define output schema with Pydantic 
class LegalSimplification(BaseModel):
    translation: str

parser = PydanticOutputParser(pydantic_object=LegalSimplification)

# Define the system prompt
format_instructions = parser.get_format_instructions()

system_message = SystemMessage(content=f"""You are a helpful assistant that translates complex legal terms into plain and understandable language.
Respond only in this format: {format_instructions}
Do not ask for clarification. Always use the given legal input.""")

# Define few-shot examples
examples = [
    {
        "input": "Notwithstanding any provision to the contrary herein, the indemnitor agrees to indemnify, defend, and hold harmless the indemnitee...",
        "output": "One party agrees to cover any costs, claims, or damages that happen because of their actions..."
    },
    {
        "input": "This agreement shall be binding upon and inure to the benefit of the parties...",
        "output": "This agreement applies to both parties and their future representatives..."
    }
]

few_shot_messages = []
for ex in examples:
    few_shot_messages.append(HumanMessage(content=ex["input"]))
    few_shot_messages.append(AIMessage(content=f'{{"translation": "{ex["output"]}"}}'))

# Define the legal input text
legal_text = (
    "Any waiver of any term or condition of this agreement shall not be deemed a continuing waiver of such term "
    "or condition, nor shall it be considered a waiver of any other term or condition hereof. No failure or delay "
    "by either party in exercising any right, power, or privilege under this agreement shall operate as a waiver "
    "thereof, nor shall any single or partial exercise preclude any other or further exercise thereof or the "
    "exercise of any other right, power, or privilege."
)

user_message = HumanMessage(content=legal_text)

# Build full message list
messages = [system_message] + few_shot_messages + [user_message]

# Sanity check the prompt
print("\n\n===== Prompt Sent to Model =====")
for m in messages:
    print(f"{m.type.upper()}: {m.content}\n")

# Create HF text-generation pipeline with lower temperature for parsing accuracy
# (temperature=0.6 here vs 0.7 in initial setup for more deterministic JSON output)
hf_pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    do_sample=True,
    temperature=0.6,
    top_p=0.9,
    return_full_text=False,
    eos_token_id=tokenizer.eos_token_id,
    skip_special_tokens=True,
)

# Wrap in LangChain-compatible Chat LLM
wrapped_llm = HuggingFacePipeline(pipeline=hf_pipe)
llm_chat = ChatHuggingFace(llm=wrapped_llm)

# Generate model output
raw_output = llm_chat.invoke(messages)

print(raw_output.content)

# Extract valid JSON from output

def extract_first_json(text):
    match = re.search(r'\{.*?\}', text, re.DOTALL)
    return match.group(0) if match else text.strip()

try:
    output_text = raw_output.content if isinstance(raw_output, AIMessage) else raw_output
    clean_output = extract_first_json(output_text)
    result = parser.parse(clean_output)

    print("\nSimplified translation:")
    print(result.translation)

    # Combine original and simplified
    entry = {
        "legal_text": legal_text,
        "translation": result.translation
    }

    # Load existing data if file exists
    data = []
    output_file = "simplified_output.json"
    if os.path.exists(output_file):
        with open(output_file, "r") as f:
            try:
                data = json.load(f)
                if not isinstance(data, list):
                    print("Warning: existing file is not a list. Overwriting.")
                    data = []
            except json.JSONDecodeError:
                data = []

    # Append new entry
    data.append(entry)

    # Write back to file
    with open(output_file, "w") as f:
        json.dump(data, f, indent=2)

    print(f"\nAppended to {output_file}")

except Exception as e:
    print("\nCould not parse output.")
    print("Raw output:", raw_output)
    print(e)

# === Unload Ollama Model & Shutdown Kernel ===
# Unloads the model from GPU memory before shutting down

try:
    import ollama
    print(f"Unloading Ollama model: {OLLAMA_LLM_MODEL}")
    ollama.generate(model=OLLAMA_LLM_MODEL, prompt="", keep_alive=0)
    print("Model unloaded from GPU memory")
except Exception as e:
    print(f"Model unload skipped: {e}")

# Shut down the kernel to fully release resources
import IPython
app = IPython.Application.instance()
app.kernel.do_shutdown(restart=False)