Upload 2 files

handler.py

@@ -0,0 +1,61 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+import json
+
+# Configuration
+BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
+ADAPTER = "tans37/mistral-query-router"
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Load the base model and tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+        
+        # Load base model in half precision for efficiency
+        base_model = AutoModelForCausalLM.from_pretrained(
+            BASE_MODEL,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        
+        # Load the Peft adapter
+        # 'path' is the directory where the handler is located (the adapter repo)
+        self.model = PeftModel.from_pretrained(base_model, path)
+        self.model.eval()
+        
+        print(f"[Handler] Loaded LoRA adapter from {path} onto {BASE_MODEL}")
+
+    def __call__(self, data):
+        """
+        Args:
+            data (:obj: `dict`):
+                subset of the request body with the following keys:
+                - `inputs`: the prompt to be processed
+                - `parameters`: optional generation parameters
+        """
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", {
+            "max_new_tokens": 128,
+            "temperature": 0.1,
+            "top_p": 0.9,
+            "do_sample": False
+        })
+
+        # Tokenize
+        inputs = self.tokenizer(inputs, return_tensors="pt").to(self.model.device)
+
+        # Generate
+        with torch.no_grad():
+            output_tokens = self.model.generate(
+                **inputs,
+                **parameters
+            )
+
+        # Decode
+        # We only want the new tokens, so we slice the output
+        input_len = inputs["input_ids"].shape[1]
+        new_tokens = output_tokens[0][input_len:]
+        prediction = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+
+        return [{"generated_text": prediction}]

requirements.txt

@@ -1,3 +1,5 @@
-transformers>=4.48.0
-peft>=0.14.0
-tokenizers>=0.21.0
+transformers>=4.40.0
+peft>=0.10.0
+torch>=2.2.0
+accelerate>=0.29.0
+bitsandbytes>=0.43.0