inference.py

# %%
from email.policy import default

from llama_cpp import Llama

from constrerl.annotator import (
    AnnotatedArticle,
    AnnotatorHelper,
    AnnotationTypes,
    Metadata,
    SpacyAnnotator,
)
from constrerl.sentences import Sentence, BERTAnnotator
from constrerl.utils import prepare_for_eval
from constrerl.beam_search.beam_search import BeamSearchConfig

# %%
import argparse
import json
from pathlib import Path

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-provider", type=str, default="llama")
    parser.add_argument(
        "--model-spec", type=str, default="quants/llama-3-2-1B-instruct-lora.gguf"
    )
    parser.add_argument(
        "--data-path", type=str, default="data/Annotations/prepared_dev_train.json"
    )
    parser.add_argument(
        "--eval-path", type=str, default="data/Articles/json_format/articles_test.json"
    )
    parser.add_argument("--out-path", type=str, default="data/results_test")
    parser.add_argument("--out-file", type=str, default="test_out.json")
    parser.add_argument("--type", type=str, default="entities")
    parser.add_argument("--top-k", type=int, default=5)
    parser.add_argument("--gen-tokens", type=int, default=512)
    parser.add_argument("--ctx", type=int, default=8196)
    parser.add_argument("--add-rag", default=False, action="store_true")
    parser.add_argument("--add-naive", default=False, action="store_true")
    parser.add_argument("--naive-filter", default=False, action="store_true")
    parser.add_argument("--naive-only", default=False, action="store_true")
    parser.add_argument("--use-ne-finetuned", default=False, action="store_true")
    parser.add_argument(
        "--beam-search", default="none", choices=["end", "shallow", "none"]
    )
    args = parser.parse_args()
    print("Starting with", args)
    model: Llama = None
    match args.model_provider:
        # case "openai":
        #     # llm = init_chat_model("ft:gpt-4o-mini-2024-07-18:tu-graz-hereditary:gutbrain-ie-finetune:B5qr9cGV", model_provider="openai")
        #     llm = init_chat_model("gpt-4o-mini-2024-07-18", model_provider="openai")

        case "llama":
            if args.model_spec.endswith(".gguf"):
                model_path = args.model_spec  # "quants/llama-3-2-1B-instruct-lora.gguf"
                model = Llama(
                    model_path,
                    n_gpu_layers=-1,
                    n_ctx=args.ctx,
                    temperature=0.1,
                    logits_all=True,
                    verbose=False,
                    # draft_model=LlamaPromptLookupDecoding(num_pred_tokens=10),
                )
            else:
                model = Llama.from_pretrained(
                    args.model_spec,
                    filename="*.Q8_0.gguf",
                    n_gpu_layers=-1,
                    n_ctx=args.ctx,
                    temperature=0.1,
                    logits_all=True,
                    verbose=False,
                )
        case "naive":
            print("Using naive annotator, no model will be loaded")
        # default:
        case _:
            print("Unknown model provider", args.model_provider)
    data_path = args.data_path
    out_path = Path(args.out_path)
    out_path.mkdir(parents=True, exist_ok=True)
    out_path = out_path / args.out_file

    beam_search = None
    gen_tokens = args.gen_tokens
    match args.beam_search:
        case "end":
            beam_search = BeamSearchConfig(k_progress=[4, 1, -1])
            gen_tokens = 2
        case "shallow":
            beam_search = BeamSearchConfig(
                top_k=2,
                max_depth=3,
                skip_tokens=2,
            )
            gen_tokens = 32
    annotator = AnnotatorHelper(
        model=model,
        gen_tokens=gen_tokens,
        add_rag=args.add_rag,
        naive_annotations=args.add_naive,
        naive_only=args.naive_only,
        naive_filter=args.naive_filter,
        top_k=args.top_k,
        beam_search=beam_search,
        ne_extractor=BERTAnnotator() if args.use_ne_finetuned else SpacyAnnotator(),
    )
    print("Loading articles from", data_path)
    annotator.load_articles_from_path(Path(data_path))
    print("-->> Loaded articles:", len(annotator.loaded_articles))
    print("Loading evaluation set from", args.eval_path)
    with open(args.eval_path, "r") as f:
        eval_set = json.load(f)
    eval_set = {
        id: Metadata.model_validate(article) for id, article in eval_set.items()
    }
    print("-->> Loaded eval set articles:", len(eval_set))

    # %%

    annotations_types = (
        [AnnotationTypes.from_str(args.type)]
        if args.type in ["entities", "relations"]
        else [AnnotationTypes.ENTITY, AnnotationTypes.RELATION]
    )
    print("Annotating with types", annotations_types)

    def save_cb(id, article, annotated_articles):
        annotator.add_concept_uris(annotated_articles)
        output_data = prepare_for_eval(annotations)
        with open(out_path, "w") as f:
            json.dump(output_data, f)

    annotations: dict[str, AnnotatedArticle] = annotator.annotate(
        {id: article for id, article in list(eval_set.items())},
        annotate=annotations_types,
        cb=save_cb,
    )

    output_data = prepare_for_eval(annotations)
    # %%
    with open(out_path, "w") as f:
        json.dump(output_data, f)
    # %%
    print("Done")