inference_example.py

"""Baguette-Software-Dataset — two-step inference example (vLLM).

Step 1 (validated):   paragraph  -> dataset / software mentions
Step 2 (functional):  mentions   -> article-level record + PLOS-OSI summary

    pip install vllm
    python inference_example.py
"""
import json
from vllm import LLM, SamplingParams

MODEL = "PleIAs/Baguette-Software-Dataset"
llm = LLM(model=MODEL, max_model_len=8192, dtype="bfloat16")
sp = SamplingParams(temperature=0.0, max_tokens=1024, stop=["<|im_end|>"])

# Exact training templates — match them.
extract = lambda para:     f"<|im_start|>user\n<text>{para}</text><|im_end|>\n<|im_start|>assistant\n"
analyze = lambda mentions: f"<|im_start|>user\n<mentions>{mentions}</mentions><|im_end|>\n<|im_start|>assistant\n"


def parse(text):
    s = text.strip(); a, b = s.find("{"), s.rfind("}")
    try: return json.loads(s[a:b + 1])
    except Exception: return {}


# A paper as an ordered list of paragraphs.
paragraphs = [
    "For the data acquisition, a T420 FLIR thermal camera with a 0.1 degree C thermal sensitivity was "
    "used. The data was acquired using the ResearchIR MAX 4.0 software by connecting the camera to a "
    "computer.",
    "Results and discussion. Figure 9 shows the infrared images obtained on the slabs containing either "
    "sound or corroded sensors.",
    "As can be seen, the intensity of the hot spot is higher for the sound sensor than the corroded one.",
]

# ---- STEP 1: per-paragraph extraction (the validated step) ----
outs = llm.generate([extract(p) for p in paragraphs], sp)
per_para = [parse(o.outputs[0].text) for o in outs]
print("=== STEP 1 — per-paragraph extraction ===")
for i, d in enumerate(per_para):
    if d.get("datasets") or d.get("software"):
        print(f"p{i}:", json.dumps(d, ensure_ascii=False))

# ---- aggregate + dedup mentions across paragraphs ----
ds, sw, seen = [], [], set()
for d in per_para:
    for x in d.get("datasets") or []:
        k = (x.get("name") or "").lower()
        if k and k not in seen: seen.add(k); ds.append(x)
    for x in d.get("software") or []:
        k = "sw:" + (x.get("name") or "").lower()
        if k and k not in seen: seen.add(k); sw.append(x)
mentions = json.dumps({"datasets": ds, "software": sw}, ensure_ascii=False)

# ---- STEP 2: article-level record + PLOS-OSI summary (indicative) ----
record = parse(llm.generate([analyze(mentions)], sp)[0].outputs[0].text)
print("\n=== STEP 2 — article record (indicative) ===")
print(json.dumps(record, ensure_ascii=False, indent=1))