#!/usr/bin/env python3
"""
Inference script for N1 Entity Extraction model.
"""

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import json

def extract_entities(text, max_new_tokens=1000):
    """Extract entities from text using the N1 model."""
    model_path = "/data/models/n1_entity_extraction_r512_dropout_merged"
    
    # Load model and tokenizer
    print("Loading N1 model...")
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True
    )
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    # Create prompt for entity extraction
    messages = [
        {"role": "system", "content": "You are an entity extraction assistant. Extract entities from the given text and output them in JSON format."},
        {"role": "user", "content": f"Extract entities from the following text:\n\n{text}"},
    ]
    
    # Apply chat template
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    
    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=16384)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
    # Generate
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.1,  # Low temperature for structured output
            do_sample=True,
            top_p=0.95,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    
    # Decode and extract assistant response
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Extract just the assistant's response
    assistant_marker = "assistant:"
    if assistant_marker in full_response:
        response = full_response.split(assistant_marker)[-1].strip()
    else:
        response = full_response
    
    # Try to parse as JSON
    try:
        entities = json.loads(response)
        return entities
    except:
        # Return raw response if not valid JSON
        return response

if __name__ == "__main__":
    # Example usage
    test_text = """
    Apple Inc. announced that Tim Cook will meet with President Biden 
    at the White House next Tuesday to discuss technology policy.
    """
    
    print("Input text:")
    print(test_text)
    print("\nExtracting entities...")
    
    entities = extract_entities(test_text)
    print("\nExtracted entities:")
    print(json.dumps(entities, indent=2) if isinstance(entities, dict) else entities)