LFM2.5-forestWHY

from transformers import AutoProcessor, AutoModelForImageTextToText
from peft import PeftModel
import torch

# ---------------------------------------------------------------
# 1. Load base model + LoRA adapter
# ---------------------------------------------------------------
BASE_MODEL = "Siddharth63/LFM2.5-forestWHY"

processor = AutoProcessor.from_pretrained(BASE_MODEL)

model = AutoModelForImageTextToText.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    dtype=torch.bfloat16,
)

# Optional: merge adapter into base for faster inference (uses more VRAM during merge)
# model = model.merge_and_unload()

# ---------------------------------------------------------------
# 2. ForestWHY prompt constants (copied from training notebook)
# ---------------------------------------------------------------
PANEL_KEYS = [
    "img_rgb_before", "img_rgb_after",
    "img_nir_false_color_before", "img_nir_false_color_after",
    "img_swir_composite_before", "img_swir_composite_after",
    "img_delta_ndvi", "img_delta_nbr",
    "img_attention_multi", "img_embedding_change",
    "img_cropa_roads", "img_delta_attn_role",
    "img_head_disagreement", "img_pca_semantic",
]

SYSTEM_PROMPT = """You are an expert remote sensing analyst and tropical ecologist
specializing in forest cover change detection from Sentinel-2 satellite imagery.
You have access to both raw spectral data AND outputs from a trained I-JEPA
Vision Transformer encoder (ViT-L/8, 24 layers, trained on 1.57M Sentinel-2 patches).

For each observation you receive 14 image panels:
SPECTRAL (1-8): RGB before/after, NIR before/after, SWIR before/after, DELTA_NDVI, DELTA_NBR
JEPA ENCODER (9-14): Multi-scale attention, Embedding change, CroPA roads,
                     Delta attention role, Head disagreement, PCA semantic clusters

Critical rules:
- Provide detailed 10-step reasoning citing specific panel numbers
- JEPA panels supersede spectral panels for ambiguous cases
- Panel 10 embedding change is more reliable than DELTA_NDVI for degradation detection
- Panel 11 CroPA road presence is the strongest predictor of continued deforestation
- Panel 14 PCA cluster split is ground truth for genuine land cover transition
- Write detailed prose for each step - minimum 4 sentences per step"""

def build_user_text(sample):
    return (
        f"Analyze this Sentinel-2 satellite observation showing land cover change.\n\n"
        f"Location: {sample.get('region','unknown')}, "
        f"{sample.get('lat',0):.3f} N {sample.get('lon',0):.3f} E\n"
        f"Period: {sample.get('year_before',2021)} to {sample.get('year_after',2025)} "
        f"({sample.get('year_gap',4)}-year span) | "
        f"Patch: 64x64 px at 19m/px approx 1.2km x 1.2km\n"
        f"Measured: DELTA_NDVI={sample.get('delta_ndvi',0):+.3f}  "
        f"DELTA_NBR={sample.get('delta_nbr',0):+.3f}\n\n"
        f"SPECTRAL PANELS (images 1-8):\n"
        f"  1. RGB Before ({sample.get('year_before',2021)})\n"
        f"  2. RGB After  ({sample.get('year_after',2025)})\n"
        f"  3. NIR False Color Before\n"
        f"  4. NIR False Color After\n"
        f"  5. SWIR Composite Before\n"
        f"  6. SWIR Composite After\n"
        f"  7. DELTA_NDVI Change Map\n"
        f"  8. DELTA_NBR Change Map\n\n"
        f"JEPA ENCODER PANELS (images 9-14) - I-JEPA ViT-L trained on 1.57M S2 patches:\n"
        f"  9.  Multi-Scale Attention (R=fine/roads G=mid/fields B=landscape)\n"
        f"  10. Embedding Change Map (cosine distance before to after per patch)\n"
        f"  11. CroPA Road Map (cross-patch token correlation for linear structures)\n"
        f"  12. Delta Attention Role (red=more salient after, blue=less salient)\n"
        f"  13. Head Disagreement Map (hot=ambiguous, cool=clear semantic signal)\n"
        f"  14. PCA Semantic Clusters (left=before, right=after)\n\n"
        f"Follow the 10-step reasoning protocol. Use JEPA panels to go beyond "
        f"what spectral indices alone can detect."
    )

# ---------------------------------------------------------------
# 3. Run inference on a dataset row (or your own 14 PIL images + metadata dict)
# ---------------------------------------------------------------
from datasets import load_dataset
dataset = load_dataset("Siddharth63/forestwhy-training-v1", split="train")
sample = dataset[0]

images = [sample[k] for k in PANEL_KEYS if sample.get(k) is not None]

conversation = [
    {"role": "system",
     "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
    {"role": "user",
     "content": [{"type": "image", "image": img} for img in images]
                + [{"type": "text", "text": build_user_text(sample)}]},
]

inputs = processor.apply_chat_template(
    conversation,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
    tokenize=True,
).to(model.device)

with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=4096,
        do_sample=True,
        temperature=1.0,
        min_p=0.1,
        use_cache=True,
    )

# Decode only the newly generated tokens
gen = outputs[:, inputs["input_ids"].shape[1]:]
print(processor.batch_decode(gen, skip_special_tokens=True)[0])