ner_inference.py

#!/usr/bin/env python3
"""
NER-Small Inference Client

A Python client for running inference with the Minibase-NER-Small model.
Handles named entity recognition requests to the local llama.cpp server.
"""

import requests
import json
from typing import Optional, Dict, Any, Tuple, List
import time
import re


class NERClient:
    """
    Client for the NER-Small named entity recognition model.

    This client communicates with a local llama.cpp server running the
    Minibase-NER-Small model for named entity recognition tasks.
    """

    def __init__(self, base_url: str = "http://127.0.0.1:8000", timeout: int = 30):
        """
        Initialize the NER client.

        Args:
            base_url: Base URL of the llama.cpp server
            timeout: Request timeout in seconds
        """
        self.base_url = base_url.rstrip('/')
        self.timeout = timeout
        self.default_instruction = "Extract all named entities from the following text. List them as 1. Entity, 2. Entity, etc."

    def _make_request(self, prompt: str, max_tokens: int = 512,
                     temperature: float = 0.1) -> Tuple[str, float]:
        """
        Make a completion request to the model.

        Args:
            prompt: The input prompt
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature

        Returns:
            Tuple of (response_text, latency_ms)
        """
        payload = {
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature
        }

        headers = {'Content-Type': 'application/json'}

        start_time = time.time()
        try:
            response = requests.post(
                f"{self.base_url}/completion",
                json=payload,
                headers=headers,
                timeout=self.timeout
            )

            latency = (time.time() - start_time) * 1000  # Convert to milliseconds

            if response.status_code == 200:
                result = response.json()
                return result.get('content', ''), latency
            else:
                return f"Error: HTTP {response.status_code}", latency

        except requests.exceptions.RequestException as e:
            latency = (time.time() - start_time) * 1000
            return f"Error: {e}", latency

    def extract_entities(self, text: str, instruction: Optional[str] = None,
                        max_tokens: int = 512, temperature: float = 0.1) -> List[Dict[str, Any]]:
        """
        Extract named entities from text.

        Args:
            text: Input text to analyze
            instruction: Custom instruction (uses default if None)
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature

        Returns:
            List of entity dictionaries with text and metadata
        """
        if instruction is None:
            instruction = self.default_instruction

        prompt = f"{instruction}\n\nInput: {text}\n\nResponse: "

        response_text, latency = self._make_request(prompt, max_tokens, temperature)

        if response_text.startswith("Error"):
            return []

        # Parse the numbered list response
        entities = self._parse_entity_response(response_text)

        # Add metadata to each entity
        for entity in entities:
            entity.update({
                'confidence': 1.0,  # Placeholder - model doesn't provide confidence
                'latency_ms': latency
            })

        return entities

    def extract_entities_batch(self, texts: List[str], instruction: Optional[str] = None,
                              max_tokens: int = 512, temperature: float = 0.1) -> List[List[Dict[str, Any]]]:
        """
        Extract named entities from multiple texts.

        Args:
            texts: List of input texts to analyze
            instruction: Custom instruction (uses default if None)
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature

        Returns:
            List of entity lists, one per input text
        """
        results = []
        for text in texts:
            entities = self.extract_entities(text, instruction, max_tokens, temperature)
            results.append(entities)

        return results

    def _parse_entity_response(self, response_text: str) -> List[Dict[str, Any]]:
        """
        Parse the model's numbered list response into structured entities.

        Args:
            response_text: Raw model response

        Returns:
            List of entity dictionaries
        """
        entities = []

        # Clean up the response
        response_text = response_text.strip()

        # Split by lines and process each line
        lines = response_text.split('\n')

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # Try to extract entity names from numbered list format
            # Pattern 1: "1. Entity Name" or "1. Entity Name - Description"
            numbered_match = re.match(r'^\d+\.\s*(.+?)(?:\s*-\s*.+)?$', line)
            if numbered_match:
                entity_text = numbered_match.group(1).strip()
                # Remove any trailing punctuation
                entity_text = re.sub(r'[.,;:!?]$', '', entity_text).strip()
                # Skip very short entities or generic terms
                if entity_text and len(entity_text) > 1 and not entity_text.lower() in ['the', 'and', 'or', 'but', 'for', 'with']:
                    entities.append({
                        'text': entity_text,
                        'type': 'ENTITY',  # Model doesn't specify types
                        'start': 0,  # Position information not available
                        'end': 0
                    })

        return entities

    def health_check(self) -> bool:
        """
        Check if the model server is healthy and responding.

        Returns:
            True if server is healthy, False otherwise
        """
        try:
            response = requests.get(f"{self.base_url}/health", timeout=5)
            return response.status_code == 200
        except:
            return False

    def get_model_info(self) -> Optional[Dict[str, Any]]:
        """
        Get information about the loaded model.

        Returns:
            Model information dictionary or None if unavailable
        """
        try:
            response = requests.get(f"{self.base_url}/v1/models", timeout=5)
            if response.status_code == 200:
                return response.json()
        except:
            pass
        return None


def main():
    """
    Command-line interface for NER inference.
    """
    import argparse

    parser = argparse.ArgumentParser(description='NER-Small Inference Client')
    parser.add_argument('text', help='Text to analyze for named entities')
    parser.add_argument('--url', default='http://127.0.0.1:8000',
                       help='Model server URL (default: http://127.0.0.1:8000)')
    parser.add_argument('--max-tokens', type=int, default=512,
                       help='Maximum tokens to generate (default: 512)')
    parser.add_argument('--temperature', type=float, default=0.1,
                       help='Sampling temperature (default: 0.1)')

    args = parser.parse_args()

    # Initialize client
    client = NERClient(args.url)

    # Check server health
    if not client.health_check():
        print(f"❌ Error: Cannot connect to model server at {args.url}")
        print("Make sure the llama.cpp server is running with the NER-Small model.")
        return 1

    # Extract entities
    entities = client.extract_entities(
        args.text,
        max_tokens=args.max_tokens,
        temperature=args.temperature
    )

    # Display results
    print(f"📝 Input Text: {args.text}")
    print(f"🎯 Found {len(entities)} entities:")
    print()

    if entities:
        for i, entity in enumerate(entities, 1):
            print(f"{i}. {entity['text']} (Type: {entity['type']})")
    else:
        print("No entities found.")

    return 0


if __name__ == "__main__":
    exit(main())