Upload 2 files

handler.py

@@ -0,0 +1,58 @@
+# 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)}

requirements.txt

@@ -0,0 +1,4 @@
+transformers>=4.42.0
+peft>=0.11.1
+accelerate>=0.29.3
+bitsandbytes==0.43.2