Initial commit: Add FLAN-T5 custom handler

handler.py

@@ -0,0 +1,50 @@
+from typing import Dict, List
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+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)
+        
+        # Enable evaluation mode
+        self.model.eval()
+
+    def __call__(self, data: Dict) -> List[Dict]:
+        # Get input text
+        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"
+        )
+        
+        # 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 

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0.0
+transformers>=4.30.0
+huggingface-hub>=0.19.0 

test_handler.py

@@ -0,0 +1,23 @@
+from handler import EndpointHandler
+
+def test_flan_t5():
+    # Initialize 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 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)
+
+if __name__ == "__main__":
+    test_flan_t5()