granite.py

"""
granite.py — IBM Granite advisory generation for AgroSense Lite.

Uses ibm-granite/granite-3.3-2b-instruct (via Hugging Face) to generate
plain-language crop disease advice grounded on structured disease metadata.

The model is lazy-loaded on first use so importing this module has no cost.
Falls back to the static advisory string if the model is unavailable.
"""

__all__ = ["generate_advice"]

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"

_tokenizer = None
_model = None


def _load_granite():
    """Download and initialise Granite on first use."""
    global _tokenizer, _model
    if _model is not None:
        return

    print("Loading IBM Granite model (first run — this may take a few minutes)...")
    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    _model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        dtype=torch.float32,   # CPU-safe; switch to bfloat16 if GPU available
        device_map="cpu",
    )
    _model.eval()
    print("Granite model ready.")


def generate_advice(condition: str, severity: str, static_advice: str, local_note: str = "") -> str:
    """
    Generate plain-language crop disease advice using IBM Granite.

    Args:
        condition:     Human-readable disease name, e.g. "Late Blight"
        severity:      Severity level: None / Moderate / High / Critical
        static_advice: Fallback advice string from the static knowledge base
        local_note:    Optional Nepal-specific context to ground the response

    Returns:
        A 2–3 sentence plain-language advisory string generated by Granite.
        Falls back to static_advice if generation fails.
    """
    try:
        _load_granite()
    except Exception as e:
        print(f"[Granite] Model load failed, using static advice. ({e})")
        return static_advice

    grounding = f"Disease: {condition}\nSeverity: {severity}\nKnown advice: {static_advice}"
    if local_note:
        grounding += f"\nLocal context (Nepal): {local_note}"

    messages = [
        {
            "role": "system",
            "content": (
                "You are an agricultural advisory assistant helping smallholder farmers in Nepal. "
                "Given structured information about a crop disease, write 2–3 sentences of clear, "
                "actionable, plain-language advice a farmer can act on immediately. "
                "Do not use bullet points or headers. Do not repeat the disease name unnecessarily. "
                "If severity is Critical, make the urgency clear."
            ),
        },
        {
            "role": "user",
            "content": grounding,
        },
    ]

    try:
        inputs = _tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
        )
        input_ids = inputs if isinstance(inputs, torch.Tensor) else inputs["input_ids"]

        with torch.no_grad():
            output_ids = _model.generate(
                input_ids,
                max_new_tokens=120,
                do_sample=False,          # deterministic — consistent advice
                pad_token_id=_tokenizer.eos_token_id,
            )

        # Decode only the newly generated tokens
        new_tokens = output_ids[0][input_ids.shape[-1]:]
        advice = _tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
        return advice if advice else static_advice

    except Exception as e:
        print(f"[Granite] Generation failed, using static advice. ({e})")
        return static_advice