Upload 2 files

e5a785e verified over 1 year ago

2.33 kB

# handler.py
# Hugging Face Inference Endpoint custom handler — April 2025 edition
from pathlib import Path
from typing import Dict, Any

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

_BASE_MODEL = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"   # 4‑bit quantised base

class EndpointHandler:
    """
    Loads the 8 B LLama‑3.1 base in 4‑bit and stitches the PEFT adapter
    found in the repository root onto it. Supports standard text‑gen kwargs.
    """

    def __init__(self, path: str = "."):
        repo = Path(path)

        # 1️⃣  Tokeniser
        self.tokenizer = AutoTokenizer.from_pretrained(
            repo if (repo / "tokenizer_config.json").exists() else _BASE_MODEL,
            padding_side="left",
            trust_remote_code=True,
        )
        self.tokenizer.pad_token = self.tokenizer.eos_token

        # 2️⃣  Base model in 4‑bit
        self.model = AutoModelForCausalLM.from_pretrained(
            _BASE_MODEL,
            load_in_4bit=True,                     # bitsandbytes
            device_map="auto",
            torch_dtype=torch.float16,
            trust_remote_code=True,
        )

        # 3️⃣  Attach LoRA / QLoRA adapter if present
        if (repo / "adapter_config.json").exists():
            self.model = PeftModel.from_pretrained(self.model, repo, is_trainable=False)
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
        prompt    = data.get("inputs") or data                       # raw string or nested JSON
        gen_cfg   = data.get("parameters", {})
        tok_in    = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)

        with torch.inference_mode():
            out = self.model.generate(
                **tok_in,
                max_new_tokens = gen_cfg.get("max_new_tokens", 256),
                temperature    = gen_cfg.get("temperature", 0.7),
                top_p          = gen_cfg.get("top_p", 0.9),
                do_sample      = gen_cfg.get("do_sample", True),
                repetition_penalty = gen_cfg.get("repetition_penalty", 1.1),
            )

        return {"generated_text": self.tokenizer.decode(out[0], skip_special_tokens=True)}

	# handler.py
	# Hugging Face Inference Endpoint custom handler — April 2025 edition
	from pathlib import Path
	from typing import Dict, Any

	import torch
	from transformers import AutoTokenizer, AutoModelForCausalLM
	from peft import PeftModel

	_BASE_MODEL = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit" # 4‑bit quantised base

	class EndpointHandler:
	"""
	Loads the 8 B LLama‑3.1 base in 4‑bit and stitches the PEFT adapter
	found in the repository root onto it. Supports standard text‑gen kwargs.
	"""

	def __init__(self, path: str = "."):
	repo = Path(path)

	# 1️⃣ Tokeniser
	self.tokenizer = AutoTokenizer.from_pretrained(
	repo if (repo / "tokenizer_config.json").exists() else _BASE_MODEL,
	padding_side="left",
	trust_remote_code=True,
	)
	self.tokenizer.pad_token = self.tokenizer.eos_token

	# 2️⃣ Base model in 4‑bit
	self.model = AutoModelForCausalLM.from_pretrained(
	_BASE_MODEL,
	load_in_4bit=True, # bitsandbytes
	device_map="auto",
	torch_dtype=torch.float16,
	trust_remote_code=True,
	)

	# 3️⃣ Attach LoRA / QLoRA adapter if present
	if (repo / "adapter_config.json").exists():
	self.model = PeftModel.from_pretrained(self.model, repo, is_trainable=False)
	self.model.eval()

	def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
	prompt = data.get("inputs") or data # raw string or nested JSON
	gen_cfg = data.get("parameters", {})
	tok_in = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)

	with torch.inference_mode():
	out = self.model.generate(
	**tok_in,
	max_new_tokens = gen_cfg.get("max_new_tokens", 256),
	temperature = gen_cfg.get("temperature", 0.7),
	top_p = gen_cfg.get("top_p", 0.9),
	do_sample = gen_cfg.get("do_sample", True),
	repetition_penalty = gen_cfg.get("repetition_penalty", 1.1),
	)

	return {"generated_text": self.tokenizer.decode(out[0], skip_special_tokens=True)}