Update handler.py

handler.py

@@ -1,101 +1,69 @@
-from typing import Any, Dict
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
+# handler.py
+import json
+from typing import List, Dict, Any
+from transformers import pipeline
 
-class EndpointHandler:
-    def __init__(self, path=""):
-        # CORRECTED: Use gemma-3-4b instead of gemma-2-2b
-        #unsloth/gemma-3-4b-it-unsloth-bnb-4bit
-        model_id = "unsloth/gemma-3-4b-it-unsloth" 
-        
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            load_in_4bit=True,
-            device_map="auto"
-        )
-        self.model = PeftModel.from_pretrained(model, path)
-    
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        # Support OpenAI Chat Completion format
-        if "messages" in data:
-            messages = data["messages"]
-            
-            # Extract system message if exists
-            system_msg = next((m["content"] for m in messages 
-                              if m["role"] == "system"), None)
-            
-            # Extract user message (last user message)
-            user_msg = next((m["content"] for m in reversed(messages) 
-                            if m["role"] == "user"), "")
-            
-            # Format prompt with system message if provided
-            if system_msg:
-                prompt = f"{system_msg}\n\nUser: {user_msg}\nAssistant:"
-            else:
-                prompt = user_msg
-        
-        # Support Text Generation format (backward compatible)
-        else:
-            prompt = data.get("inputs", "")
-        
-        # Extract generation parameters
-        params = data.get("parameters", {})
-        max_new_tokens = params.get("max_new_tokens", 128)
-        temperature = params.get("temperature", 0.7)
-        top_p = params.get("top_p", 0.9)
-        do_sample = params.get("do_sample", True)
-        
-        # Tokenize input
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-        
-        # Generate response
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=do_sample,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id
-            )
-        
-        # Decode generated text
-        generated_text = self.tokenizer.decode(
-            outputs[0], 
-            skip_special_tokens=True
-        )
-        
-        # Remove the input prompt from output
-        if prompt in generated_text:
-            generated_text = generated_text.replace(prompt, "").strip()
-        
-        # Return in appropriate format based on input type
-        if "messages" in data:
-            # OpenAI Chat Completion response format
-            return {
-                "id": "chatcmpl-pidgin",
-                "object": "chat.completion",
-                "created": int(torch.cuda.Event().record().elapsed_time(torch.cuda.Event())),
-                "model": "pidgin-gemma-3",
-                "choices": [{
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": generated_text
-                    },
-                    "finish_reason": "stop"
-                }],
-                "usage": {
-                    "prompt_tokens": len(inputs["input_ids"][0]),
-                    "completion_tokens": len(outputs[0]) - len(inputs["input_ids"][0]),
-                    "total_tokens": len(outputs[0])
-                }
+# ---- LOAD MODEL ON STARTUP ----
+pipe = pipeline(
+    "text-generation",
+    model="Ephraimmm/PIDGIN_gemma-3",
+    device_map="auto"
+)
+
+
+# ---- UTIL: convert OpenAI-like body into prompt ----
+def build_prompt(messages: List[Dict[str, str]]) -> str:
+    out = []
+    for m in messages:
+        role = m["role"]
+        content = m["content"]
+        if role == "system":
+            out.append(f"<<SYS>> {content}\n")
+        elif role == "user":
+            out.append(f"User: {content}")
+        elif role == "assistant":
+            out.append(f"Assistant: {content}")
+    out.append("Assistant:")
+    return "\n".join(out)
+
+
+# ---- MAIN HANDLER ENTRYPOINT ----
+def handle(event, context=None):
+    """AWS/SageMaker/HF style single entry."""
+    if isinstance(event, str):  # if raw body
+        event = json.loads(event)
+
+    # Expect OpenAI style body
+    messages = event["messages"]
+    max_tokens = event.get("max_tokens", 256)
+    temperature = event.get("temperature", 0.7)
+    top_p = event.get("top_p", 0.95)
+
+    prompt = build_prompt(messages)
+
+    out = pipe(
+        prompt,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+        eos_token_id=pipe.tokenizer.eos_token_id,
+    )[0]["generated_text"]
+
+    # extract only the assistant continuation
+    completion = out.replace(prompt, "").strip()
+
+    # ---- RETURN IN OPENAI CHAT FORMAT ----
+    resp = {
+        "id": "chatcmpl-local-1",
+        "object": "chat.completion",
+        "choices": [{
+            "index": 0,
+            "finish_reason": "stop",
+            "message": {
+                "role": "assistant",
+                "content": completion
             }
-        else:
-            # Text Generation response format
-            return [{
-                "generated_text": generated_text
-            }]
+        }],
+    }
+    return resp