schema-gen-llama3-8b-stage1-merged

Full merged model from Stage‑3 (weights + tokenizer).

Recommended inference (merged, no PEFT):

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

BASE = "mohdusman001/schema-gen-llama3-8b-stage1-merged"  # this merged repo

SYSTEM_PROMPT = (
    "You induce minimal JSON schemas from documents. "
    "Output strictly valid JSON with no commentary."
)

def build_user(document_text: str) -> str:
    return document_text.strip()

def generate_schema(text: str, max_new_tokens: int = 320, base_repo: str = BASE) -> str:
    tok = AutoTokenizer.from_pretrained(base_repo, use_fast=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        base_repo,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": build_user(text)},
    ]

    if hasattr(tok, "apply_chat_template"):
        prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    else:
        prompt = f"<s>[SYSTEM]\n{SYSTEM_PROMPT}\n[/SYSTEM]\n[USER]\n{text}\n[/USER]\n"

    inputs = tok(prompt, return_tensors="pt")
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.0,
            do_sample=False,
            eos_token_id=tok.eos_token_id,
            pad_token_id=tok.pad_token_id,
        )

    result = tok.decode(out[0], skip_special_tokens=True).strip()
    return result

Artifacts (if uploaded):

• Eval: eval/final_eval.json
• Samples: samples/generations.jsonl

Primary LoRA adapter lives at: https://huggingface.co/mohdusman001/schema-gen-llama3-8b-stage1-lora