Update handler.py

handler.py

@@ -1,117 +1,5 @@
 # handler.py
 
-# from typing import Any, Dict
-# import torch
-# from transformers import AutoModelForCausalLM, AutoTokenizer
-# from huggingface_inference_toolkit.logging import logger
-
-# model_dir = 'Ephraimmm/PIDGIN_gemma-3'
-
-
-
-# class EndpointHandler:
-#     """
-#     Custom handler for Hugging Face Inference Endpoints.
-#     """
-
-#     def __init__(self, model_dir: str, **kwargs: Any):
-#         """
-#         Called once at startup. `model_dir` is the path to the downloaded model repository.
-#         You can load your model, tokenizer, etc. here.
-#         """
-#         logger.info(f"Initializing EndpointHandler with model_dir = {model_dir}")
-
-#         # load tokenizer and model from the model_dir or hub
-#         base_model_name = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit" 
-#         base_model, tokenizer = FastLanguageModel.from_pretrained(
-#         model_name=base_model_name,
-#         max_seq_length=2048,
-#         dtype=torch.float16,
-#         load_in_4bit=True,
-#         r = 8,           # Larger = higher accuracy, but might overfit
-#         lora_alpha = 8,  # Recommended alpha == r at least
-#         lora_dropout = 0,
-#         )
-        
-#         # -----------------------------
-#         # 3️⃣ Load LoRA
-#         # -----------------------------
-#         lora_repo =  "Ephraimmm/PIDGIN_gemma-3" #"Ephraimmm/pigin-gemma-3-0.2" #"Ephraimmm/Pidgin_llamma_model" 
-#         lora_model = PeftModel.from_pretrained(base_model, lora_repo, adapter_name="adapter_model")
-#         FastLanguageModel.for_inference(lora_model)
-
-#         # we assume model files are present in model_dir
-#         # try:
-#         #     self.tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=False)
-#         #     self.model = AutoModelForCausalLM.from_pretrained(model_dir)
-#         # except Exception as e:
-#         #     # fallback: try loading from hub name
-#         #     hf_model_id = "Ephraimmm/PIDGIN_gemma-3"
-#         #     logger.warn(f"Loading model from hub id {hf_model_id} because fallback. Exception: {e}")
-#         #     self.tokenizer = AutoTokenizer.from_pretrained(hf_model_id, use_fast=False)
-#         #     self.model = AutoModelForCausalLM.from_pretrained(hf_model_id)
-
-#         # put the model on GPU if available
-#         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-#         self.model = self.model.to(self.device)
-#         self.model.eval()
-
-#         # You could also load extra assets, set up config, etc.
-
-#     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-#         """
-#         Called on each API request. `data` is the parsed JSON body from the request.
-#         Must return a JSON-serializable structure.
-#         The `data` dict always has an “inputs” key (unless you change it).
-#         """
-#         logger.info(f"EndpointHandler received data: {data}")
-
-#         # Extract the prompt (the "inputs" key)
-#         prompt = data.get("inputs")
-#         if prompt is None:
-#             raise ValueError("Request JSON must contain an `inputs` field with the prompt string.")
-
-#         # optional: support generation parameters from the payload
-#         max_new_tokens = data.get("parameters", {}).get("max_new_tokens", 50)
-#         temperature = data.get("parameters", {}).get("temperature", 1.0)
-#         top_k = data.get("parameters", {}).get("top_k", 50)
-#         top_p = data.get("parameters", {}).get("top_p", 1.0)
-#         do_sample = data.get("parameters", {}).get("do_sample", True)
-
-#         # Tokenize
-#         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
-
-#         # Generate
-#         output = self.model.generate(
-#             **inputs,
-#             max_new_tokens=max_new_tokens,
-#             temperature=temperature,
-#             top_k=top_k,
-#             top_p=top_p,
-#             do_sample=do_sample,
-#         )
-
-#         # Decode
-#         full = self.tokenizer.decode(output[0], skip_special_tokens=True)
-
-#         # Optionally remove the original prompt from the returned text to send only the continuation
-#         continuation = full[len(prompt):] if full.startswith(prompt) else full
-
-#         # Prepare response in the shape expected by Inference Endpoints
-#         # The “predictions” key is often used (the default inference service expects it)
-#         return {
-#             "predictions": [full],
-#             "continuations": [continuation],
-#         }
-
-
-# # For local testing (not used in production)
-# if __name__ == "__main__":
-#     # Quick local test
-#     h = EndpointHandler(model_dir=".")
-#     resp = h({"inputs": "Once upon a time", "parameters": {"max_new_tokens": 20}})
-#     print("Response:", resp)
-
 from typing import Any, Dict, List, Union
 
 import torch
@@ -231,11 +119,11 @@ class EndpointHandler:
 
         params = data.get("parameters", {}) or {}
  
-        max_new_tokens = int(params.get("max_new_tokens", 128))
+        max_new_tokens = int(params.get("max_new_tokens", 64))
 
-        temperature    = float(params.get("temperature", 0.7))
+        temperature    = float(params.get("temperature", 1))
 
-        top_k          = int(params.get("top_k", 50))
+        top_k          = int(params.get("top_k", 64))
 
         top_p          = float(params.get("top_p", 0.95))
 
@@ -277,9 +165,9 @@ class EndpointHandler:
  
         return {
 
-            "predictions": [full_text],
+            "reply": [full_text],
 
-            "continuations": [continuation],
+            # "continuations": [continuation],
 
             "usage": {
 
@@ -302,7 +190,7 @@ if __name__ == "__main__":
 
         "inputs": "Write in Nigerian Pidgin: How you dey?",
 
-        "parameters": {"max_new_tokens": 40}
+        "parameters": {"max_new_tokens": 64}
 
     })