求助：Java Deep Java Library下Hugging Face all-MiniLM-L6-v2模型预测问题及Translator示例

1. 先确认依赖安装 #

如果还没装必要的库，执行以下命令：

pip install transformers sentence-transformers

2. 自定义Translator实现（简洁版） #

这个示例里的Translator会把输入字符串转换成句子嵌入向量，再输出你需要的float类型结果（这里以向量均值为例，你可以根据需求修改逻辑）：

from sentence_transformers import SentenceTransformer

class CustomTextTranslator:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        # 加载预训练模型
        self.model = SentenceTransformer(model_name)
    
    def translate(self, input_text: str) -> float:
        # 获取文本的嵌入向量
        embedding = self.model.encode(input_text, convert_to_tensor=False)
        # 这里把向量转换成float，示例取向量的均值
        # 你也可以改成取向量的最大值、最小值，或者某个特定索引的元素
        result = embedding.mean()
        return result

# 实例化Translator并测试
translator = CustomTextTranslator()
test_text = "Hello, this is a sample sentence."
output_float = translator.translate(test_text)
print(f"输入文本: {test_text}")
print(f"输出float值: {output_float}")

3. 原生Transformers库实现（手动处理池化） #

如果更习惯用原生transformers库，需要自己实现句子嵌入的池化步骤，示例如下：

from transformers import AutoTokenizer, AutoModel
import torch

class PipelineTranslator:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)
    
    def mean_pooling(self, model_output, attention_mask):
        # 均值池化获取句子嵌入
        token_embeddings = model_output[0]
        input_mask = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask, 1) / torch.clamp(input_mask.sum(1), min=1e-9)
    
    def translate(self, input_text: str) -> float:
        # 预处理文本
        encoded_input = self.tokenizer(input_text, padding=True, truncation=True, return_tensors='pt').to(self.device)
        # 得到模型输出
        with torch.no_grad():
            model_output = self.model(**encoded_input)
        # 生成句子嵌入
        sentence_embedding = self.mean_pooling(model_output, encoded_input['attention_mask'])
        # 转换为float结果，示例取均值
        result = sentence_embedding.cpu().numpy().mean()
        return float(result)

# 实例化并测试
pipe_translator = PipelineTranslator()
test_text = "This is another test case."
output = pipe_translator.translate(test_text)
print(f"输出float值: {output}")

关键调整点 #

如果需要的不是向量均值，直接修改translate方法里的处理逻辑即可——比如取向量的最大值embedding.max()、某个索引值embedding[0]，或者其他自定义计算方式。

内容的提问来源于stack exchange，提问作者Dinesh