Initial commit: Add FLAN-T5 custom handler

handler.py

@@ -1,50 +1,39 @@
-from typing import Dict, List
-import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
 
-class EndpointHandler():
-    def __init__(self, path=""):
-        # Load FLAN-T5 model and tokenizer
-        self.model_name = "google/flan-t5-large"
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
+class EndpointHandler:
+    def __init__(self, path="google/flan-t5-large"):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
         
-        # Enable evaluation mode
-        self.model.eval()
-
-    def __call__(self, data: Dict) -> List[Dict]:
-        # Get input text
+    def __call__(self, data):
+        """
+        Args:
+            data: (dict): A dictionary with a "inputs" key containing the text to process
+        """
         inputs = data.pop("inputs", data)
         
-        # Ensure inputs is a list
-        if isinstance(inputs, str):
-            inputs = [inputs]
-            
-        # Tokenize inputs
-        tokenized = self.tokenizer(
-            inputs,
-            padding=True,
-            truncation=True,
-            max_length=512,
-            return_tensors="pt"
-        )
+        # Parameters for text generation
+        parameters = {
+            "max_length": 512,
+            "min_length": 32,
+            "temperature": 0.9,
+            "top_p": 0.95,
+            "top_k": 50,
+            "do_sample": True,
+            "num_return_sequences": 1
+        }
+        
+        # Update parameters if provided in the request
+        parameters.update(data)
+        
+        # Tokenize the input
+        input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids
+        
+        # Generate the response
+        outputs = self.model.generate(input_ids, **parameters)
+        
+        # Decode the response
+        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Perform inference
-        with torch.no_grad():
-            outputs = self.model.generate(
-                tokenized.input_ids,
-                max_length=512,
-                min_length=50,
-                temperature=0.9,
-                top_p=0.95,
-                top_k=50,
-                do_sample=True,
-                num_return_sequences=1
-            )
-            
-        # Decode the generated responses
-        responses = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-            
-        # Format output
-        results = [{"generated_text": response} for response in responses]
-        return results 
+        return {"generated_text": generated_text} 

test_handler.py

@@ -1,23 +1,19 @@
 from handler import EndpointHandler
 
-def test_flan_t5():
-    # Initialize handler
+def test_handler():
+    # Initialize the handler
     handler = EndpointHandler()
     
-    # Test cases
-    test_inputs = [
-        "Explain quantum computing in simple terms",
-        "Translate 'Hello, how are you?' to French",
-        "Write a short story about a magical forest"
-    ]
+    # Test with a simple prompt
+    test_input = {
+        "inputs": "Explain quantum computing in simple terms"
+    }
     
-    # Test each input
-    for text in test_inputs:
-        print("\n" + "="*50)
-        print(f"Input text: {text}")
-        result = handler({"inputs": text})
-        print("Generated response:", result)
-        print("="*50)
+    # Get the response
+    response = handler(test_input)
+    
+    print("Input:", test_input["inputs"])
+    print("Output:", response["generated_text"])
 
 if __name__ == "__main__":
-    test_flan_t5() 
+    test_handler()