Update inference.py

inference.py

@@ -1,18 +1,47 @@
-# inference.py
-from tokenizers import Tokenizer
-from typing import List, Dict
+from transformers import Pipeline, PreTrainedTokenizer, AutoTokenizer
+from typing import Dict, Union, List
+import torch
 
-# 加载 tokenizer
-tokenizer = Tokenizer.from_pretrained(".")  # 从当前目录加载
+class TokenizerPipeline(Pipeline):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
 
-def tokenize(text: str) -> Dict[str, List[str]]:
-    """接收文本，返回分词结果"""
-    encoded = tokenizer.encode(text)
-    return {
-        "tokens": encoded.tokens,
-        "ids": encoded.ids
-    }
+    def _sanitize_parameters(self, **kwargs):
+        # 处理传入参数：是否解码、padding等
+        preprocess_kwargs = {}
+        if "padding" in kwargs:
+            preprocess_kwargs["padding"] = kwargs["padding"]
+        if "truncation" in kwargs:
+            preprocess_kwargs["truncation"] = kwargs["truncation"]
+        
+        postprocess_kwargs = {}
+        if "return_tokens" in kwargs:
+            postprocess_kwargs["return_tokens"] = kwargs["return_tokens"]
+        
+        return preprocess_kwargs, {}, postprocess_kwargs
 
-# 测试示例（可选）
-if __name__ == "__main__":
-    print(tokenize("Hello, this is a test."))
+    def preprocess(self, inputs, **kwargs) -> Dict:
+        # 使用Tokenizer处理输入文本
+        return self.tokenizer(inputs, return_tensors="pt", **kwargs)
+
+    def _forward(self, inputs) -> Dict:
+        # 直接返回预处理结果（无模型推理）
+        return inputs
+
+    def postprocess(self, model_outputs, **kwargs) -> Dict:
+        # 转换输出为可读格式
+        input_ids = model_outputs["input_ids"][0]
+        
+        if kwargs.get("return_tokens", True):
+            tokens = self.tokenizer.convert_ids_to_tokens(input_ids)
+            return {"tokens": tokens}
+        else:
+            return {"input_ids": input_ids.tolist()}
+
+# 关键：创建并导出pipeline实例
+tokenizer = AutoTokenizer.from_pretrained(".")
+pipeline = TokenizerPipeline(tokenizer=tokenizer)
+
+# 可选：添加类型提示供HF解析
+def get_pipeline() -> Pipeline:
+    return pipeline