updated

merged_distilgpt2/handler.py

@@ -1,11 +1,13 @@
 from typing import Dict, List, Any
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline, AutoConfig, AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification
 from sentence_transformers import SentenceTransformer
 import torch
+import os
 
 class EndpointHandler:
     def __init__(self, path=""):
-        self.task = self._determine_task(path)
+        self.path = path
+        self.task = self._determine_task()
         if self.task == "text-generation":
             self.model = AutoModelForCausalLM.from_pretrained(path)
             self.tokenizer = AutoTokenizer.from_pretrained(path)
@@ -27,45 +29,49 @@ class EndpointHandler:
         elif self.task == "sentence-embedding":
             self.model = SentenceTransformer(path)
         else:
-            raise ValueError(f"Unsupported task: {self.task}")
+            raise ValueError(f"Unsupported task: {self.task} for model at {path}")
 
-    def _determine_task(self, path):
-        model_name = path.split("/")[-1]
-        text_generation_models = [
-            "distilgpt2",
-            "fine_tuned_distilgpt2_lora",
-            "fine_tuned_gpt2",
-            "merged_distilgpt2",
-            "gpt2"
-        ]
-        text_classification_models = [
-            "emotion_classifier",
-            "emotion_model",
-            "intent_classifier",
-            "intent_fallback"
-        ]
-        embedding_models = ["intent_encoder", "sentence_transformer"]
+    def _determine_task(self):
+        # Load config to determine model_type
+        config_path = os.path.join(self.path, "config.json")
+        if not os.path.exists(config_path):
+            raise ValueError(f"config.json not found in {self.path}")
         
-        if model_name in text_generation_models:
+        config = AutoConfig.from_pretrained(self.path)
+        model_type = config.model_type if hasattr(config, "model_type") else None
+        
+        # Map model_type or model name to tasks
+        text_generation_types = ["gpt2"]
+        text_classification_types = ["bert", "distilbert", "roberta"]
+        embedding_types = ["bert"]  # Sentence-BERT models use bert model_type
+        
+        model_name = self.path.split("/")[-1]
+        if model_type in text_generation_types or model_name in ["fine_tuned_gpt2", "merged_distilgpt2"]:
             return "text-generation"
-        elif model_name in text_classification_models:
+        elif model_type in text_classification_types or model_name in ["emotion_classifier", "intent_classifier", "intent_fallback"]:
             return "text-classification"
-        elif model_name in embedding_models:
+        elif model_name in ["intent_encoder", "sentence_transformer"] or "sentence_bert_config.json" in os.listdir(self.path):
             return "sentence-embedding"
-        return "text-generation"  # Default
+        elif model_type in text_classification_types and model_name == "emotion_model":
+            # Handle emotion_model, which may be classification or generation
+            return "text-classification"  # Assume classification; adjust if needed
+        raise ValueError(f"Could not determine task for model_type: {model_type}, model_name: {model_name}")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         inputs = data.get("inputs", "")
         if not inputs:
             return [{"error": "No inputs provided"}]
 
-        if self.task == "text-generation":
-            result = self.pipeline(inputs, max_length=50, num_return_sequences=1)
-            return [{"generated_text": item["generated_text"]} for item in result]
-        elif self.task == "text-classification":
-            result = self.pipeline(inputs, return_all_scores=True)
-            return result
-        elif self.task == "sentence-embedding":
-            embeddings = self.model.encode(inputs)
-            return [{"embeddings": embeddings.tolist()}]
-        return [{"error": f"Unsupported task: {self.task}"}]
+        try:
+            if self.task == "text-generation":
+                result = self.pipeline(inputs, max_length=50, num_return_sequences=1)
+                return [{"generated_text": item["generated_text"]} for item in result]
+            elif self.task == "text-classification":
+                result = self.pipeline(inputs, return_all_scores=True)
+                return [{"label": item["label"], "score": item["score"]} for sublist in result for item in sublist]
+            elif self.task == "sentence-embedding":
+                embeddings = self.model.encode(inputs)
+                return [{"embeddings": embeddings.tolist()}]
+            return [{"error": f"Unsupported task: {self.task}"}]
+        except Exception as e:
+            return [{"error": f"Inference failed: {str(e)}"}]