Upload handler.py

handler.py

@@ -2,65 +2,70 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import os
 
-# Global model and tokenizer instances
-model = None
-tokenizer = None
+# Create the required EndpointHandler class for Hugging Face Inference Endpoints
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Initialize model and tokenizer as None - will be loaded on first request
+        self.model = None
+        self.tokenizer = None
+        self.path = path
+        print(f"Initializing EndpointHandler with path: {path}")
 
-def load_model():
-    """Initialize the model and tokenizer once"""
-    global model, tokenizer
+    def __call__(self, data, parameters=None):
+        # This will be called when the endpoint receives a request
+        return self.predict(data, parameters)
     
-    # Model repository ID
-    model_id = "EvolphTech/Wildnerve-tlm01_Hybrid_Model"
-    
-    # Get token from environment (if needed for private models)
-    hf_token = os.environ.get("HF_TOKEN")
-    
-    # Load model and tokenizer
-    model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_token)
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
-    
-    # Move to GPU if available
-    if torch.cuda.is_available():
-        model = model.to("cuda")
-    
-    print("Model and tokenizer loaded successfully!")
+    def load_model(self):
+        """Initialize the model and tokenizer once"""
+        # Model repository ID
+        model_id = "EvolphTech/Wildnerve-tlm01_Hybrid_Model"
+        
+        # Get token from environment (if needed for private models)
+        hf_token = os.environ.get("HF_TOKEN")
+        
+        # Load model and tokenizer
+        self.model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_token)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+        
+        # Move to GPU if available
+        if torch.cuda.is_available():
+            self.model = self.model.to("cuda")
+        
+        print("Model and tokenizer loaded successfully!")
 
-def predict(inputs, parameters=None):
-    """Primary inference function for the model"""
-    global model, tokenizer
-    
-    # Load model if not already loaded
-    if model is None or tokenizer is None:
-        load_model()
-    
-    # Default parameters
-    max_length = parameters.get("max_length", 100) if parameters else 100
-    temperature = parameters.get("temperature", 0.7) if parameters else 0.7
-    
-    # Process the input text
-    if isinstance(inputs, str):
-        text_input = inputs
-    elif isinstance(inputs, dict) and "inputs" in inputs:
-        text_input = inputs["inputs"]
-    else:
-        text_input = str(inputs)
-    
-    # Tokenize inputs
-    inputs = tokenizer(text_input, return_tensors="pt")
-    input_ids = inputs["input_ids"].to(model.device)
-    
-    # Generate text
-    with torch.no_grad():
-        outputs = model.generate(
-            input_ids,
-            max_length=max_length,
-            temperature=temperature,
-            do_sample=temperature > 0,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    
-    # Decode and return the generated text
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    return {"generated_text": generated_text}
+    def predict(self, inputs, parameters=None):
+        """Primary inference function for the model"""
+        # Load model if not already loaded
+        if self.model is None or self.tokenizer is None:
+            self.load_model()
+        
+        # Default parameters
+        max_length = parameters.get("max_length", 100) if parameters else 100
+        temperature = parameters.get("temperature", 0.7) if parameters else 0.7
+        
+        # Process the input text
+        if isinstance(inputs, str):
+            text_input = inputs
+        elif isinstance(inputs, dict) and "inputs" in inputs:
+            text_input = inputs["inputs"]
+        else:
+            text_input = str(inputs)
+        
+        # Tokenize inputs
+        inputs = self.tokenizer(text_input, return_tensors="pt")
+        input_ids = inputs["input_ids"].to(self.model.device)
+        
+        # Generate text
+        with torch.no_grad():
+            outputs = self.model.generate(
+                input_ids,
+                max_length=max_length,
+                temperature=temperature,
+                do_sample=temperature > 0,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+        
+        # Decode and return the generated text
+        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return {"generated_text": generated_text}