Update handler.py

handler.py

@@ -1,47 +1,43 @@
-from typing import Dict, Any
+import os
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-from peft import PeftModel, PeftConfig, prepare_model_for_kbit_training
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from peft import PeftModel, PeftConfig
 
 class EndpointHandler:
-    def __init__(self, model_dir: str):
-        # Load tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    def __init__(self, model_dir):
+        # Optional: Get Hugging Face token for gated models
+        hf_token = os.getenv("HF_TOKEN")
 
-        # Load base model with quantization
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4"
+        # Load the PEFT adapter config
+        peft_config = PeftConfig.from_pretrained(model_dir)
+
+        # Load the base model and tokenizer (add use_auth_token if model is gated)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            peft_config.base_model_name_or_path,
+            use_auth_token=hf_token
         )
         base_model = AutoModelForCausalLM.from_pretrained(
-            model_dir,
-            quantization_config=bnb_config,
-            device_map="auto",
-            trust_remote_code=True
+            peft_config.base_model_name_or_path,
+            use_auth_token=hf_token
         )
 
-        # Prepare for k-bit training (needed for PEFT)
-        base_model = prepare_model_for_kbit_training(base_model)
-
-        # Load PEFT adapter
+        # Load the adapter onto the base model
         self.model = PeftModel.from_pretrained(base_model, model_dir)
+        self.model.eval()  # Set to evaluation mode
 
-        self.model.eval()
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        inputs = data.get("inputs", "")
-        gen_kwargs = data.get("parameters", {"max_new_tokens": 256, "do_sample": True})
+        # Set device for pipeline: 0 for GPU, -1 for CPU
+        device = 0 if torch.cuda.is_available() else -1
 
-        inputs_tokenized = self.tokenizer(inputs, return_tensors="pt").to(self.model.device)
-
-        with torch.no_grad():
-            outputs = self.model.generate(
-                input_ids=inputs_tokenized["input_ids"],
-                attention_mask=inputs_tokenized["attention_mask"],
-                **gen_kwargs
-            )
+        # Create the text generation pipeline with device assignment
+        self.generator = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            device=device
+        )
 
-        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return {"generated_text": generated_text}
+    def __call__(self, inputs):
+        text = inputs.get("inputs", "")
+        # You can add more generation parameters here as needed
+        results = self.generator(text)
+        return results