Upload handler.py with huggingface_hub

handler.py

@@ -0,0 +1,82 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+def EndpointHandler(path=""):
+    """
+    Inference Endpoints에서 사용할 핸들러 함수
+    """
+    class Handler:
+        def __init__(self, path=""):
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.base_model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+            self.adapter_path = path or "./tinyllama-qa-news"
+
+            # ✅ 모델 로딩 (GPU)
+            base_model = AutoModelForCausalLM.from_pretrained(
+                self.base_model_id,
+                torch_dtype=torch.float16
+            )
+            self.model = PeftModel.from_pretrained(base_model, self.adapter_path)
+            self.model.to(self.device)
+            self.model.eval()
+
+            self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            self.tokenizer.padding_side = "right"
+
+        def generate_response(self, user_input: str):
+            messages = [
+                {
+                    "role": "system",
+                    "content": "당신은 공감적이고 이해심 깊은 감정 상담 챗봇입니다. 사용자의 감정과 생각을 이해하고 공감하며, 적절한 조언을 제공합니다."
+                },
+                {"role": "user", "content": user_input}
+            ]
+
+            prompt = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                return_tensors=None
+            )
+
+            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True)
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=256,
+                    temperature=0.8,
+                    top_p=0.95,
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    repetition_penalty=1.1,
+                    length_penalty=1.0,
+                    num_return_sequences=1,
+                    num_beams=1,
+                    early_stopping=False
+                )
+
+            decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+            if "<|assistant|>" in decoded:
+                response = decoded.split("<|assistant|>")[-1].strip()
+            else:
+                response = decoded.split(prompt)[-1].strip()
+
+            if not response:
+                response = "감정을 이해해 주셔서 감사합니다. 더 이야기하고 싶으신 게 있으신가요?"
+
+            return response
+
+        def __call__(self, data):
+            if not data or "inputs" not in data:
+                return {"error": "입력 데이터가 없습니다."}
+
+            user_input = data["inputs"]
+            response = self.generate_response(user_input)
+            return {"response": response}
+    
+    return Handler(path)