Create handler.py

handler.py

@@ -0,0 +1,78 @@
+import os
+import torch
+import numpy as np
+import joblib
+
+# Hugging Face가 제공하는 기본 핸들러 클래스를 가져옵니다.
+from text_generation_server.models.custom_handler import BaseHandler
+
+class TidePredictionHandler(BaseHandler):
+    """
+    TimeXer 모델을 위한 커스텀 핸들러
+    """
+    def __init__(self, model, tokenizer):
+        # 이 함수는 서버가 시작될 때 딱 한 번 실행됩니다.
+        # 모델 초기화, 스케일러 로딩 등 준비 작업을 여기서 합니다.
+        super().__init__(model, tokenizer)
+        
+        # 1. 스케일러를 불러와서 self.scaler에 저장합니다.
+        #    경로는 저장소 내의 실제 파일 위치와 같아야 합니다.
+        scaler_path = os.path.join(os.getcwd(), 'checkpoints', 'scaler.gz')
+        self.scaler = joblib.load(scaler_path)
+        
+        # 2. 모델을 평가 모드로 설정합니다.
+        self.model.eval()
+        
+        # 3. 모델의 설정값들을 self.model.args 처럼 접근할 수 있도록 저장해두면 편리합니다.
+        #    이 부분은 TimeXer 모델의 구조에 따라 필요 없을 수도 있습니다.
+        #    예: self.seq_len = self.model.seq_len
+
+    def __call__(self, inputs, **kwargs):
+        # 이 함수는 API 예측 요청이 올 때마다 실행됩니다.
+        # 실제 예측 로직이 들어가는 부분입니다.
+
+        # 1. 입력 데이터 파싱
+        # inputs는 보통 리스트 형태의 텍스트 또는 바이트 데이터로 들어옵니다.
+        # JSON 형식으로 입력을 받으려면 추가적인 처리가 필요할 수 있습니다.
+        # 여기서는 간단히 inputs가 숫자 리스트 문자열이라고 가정합니다.
+        # 예: "500.1, 502.3, ..., 498.7"
+        
+        # 문자열을 숫자 리스트로 변환
+        try:
+            # 입력 데이터를 파싱하는 가장 좋은 방법은 JSON을 사용하는 것입니다.
+            # 예: `json.loads(inputs[0])`
+            # 여기서는 간단한 예시를 위해 split을 사용합니다.
+            input_list = [float(i) for i in inputs[0].split(',')]
+            seq_len = 144 # 이 값은 실제 모델의 입력 길이와 일치해야 합니다.
+
+            if len(input_list) != seq_len:
+                 raise ValueError(f"Input must have {seq_len} items.")
+
+        except Exception as e:
+            # 오류 발생 시 에러 메시지를 반환합니다.
+            return {"error": f"Invalid input format: {str(e)}"}, 400
+
+        # 2. 데이터를 모델 입력 형식(Tensor)으로 변환
+        input_array = np.array(input_list).reshape(-1, 1)
+        scaled_input = self.scaler.transform(input_array)
+        input_tensor = torch.from_numpy(scaled_input).float().unsqueeze(0).to(self.model.device)
+
+        # 3. 모델 예측 실행
+        with torch.no_grad():
+            # TimeXer 모델의 forward 함수에 필요한 모든 인자를 전달해야 합니다.
+            # 예: outputs = self.model(batch_x=input_tensor, batch_x_mark=...)
+            # 이 부분은 모델의 실제 코드를 보고 채워야 합니다.
+            # 여기서는 input_tensor만 필요하다고 가정합니다.
+            outputs = self.model(input_tensor)
+
+        # 4. 예측 결과를 후처리하고 원래 스케일로 복원
+        prediction_scaled = outputs.detach().cpu().numpy().squeeze()
+        prediction = self.scaler.inverse_transform(prediction_scaled.reshape(-1, 1))
+        
+        # 5. 최종 결과를 리스트 형태로 반환
+        # Hugging Face 핸들러는 보통 텍스트나 바이트를 반환해야 합니다.
+        # 결과를 JSON 문자열로 만들어 반환하는 것이 일반적입니다.
+        import json
+        result_str = json.dumps({"prediction": prediction.flatten().tolist()})
+
+        return [result_str]