fix conflicts

Dockerfile

@@ -0,0 +1,6 @@
+# FROM quay.io/ascend/vllm-ascend:v0.10.0rc1
+FROM git.modelhub.org.cn:9443/enginex-ascend/vllm-ascend:v0.10.0rc1
+
+WORKDIR /workspace
+RUN pip install sentence-transformers
+COPY main.py test.sh dataset.json /workspace/

README.md

@@ -1,3 +1,27 @@
 # enginex-ascend-910-fe
 
 运行于【昇腾-910】系列算力卡的【特征抽取】引擎，基于 transformer 架构，支持 BGE、jina-clip 等最新流行模型
+
+## Quickstart
+
+### 构建镜像
+```bash
+docker build -t feature:v0.1 .
+```
+
+### 模型下载
+模型地址：https://modelscope.cn/models/BAAI/bge-large-zh-v1.5
+并放到目录：`/mnt/contest_ceph/zhanghao/models/BAAI/bge-large-zh-v1.5`（如更改目录，请修改后面的执行脚本中的模型路径）
+
+### 测试程序
+1. 准备输入数据集，可以参考示例`dataset.json`
+2. 在docker镜像里运行测试程序，会根据`dataset.json`内容，计算每个句子的embedding，同时计算所有句子的两两相似度，结果保存在`output.json`。
+```bash
+./run_in_docker.sh
+```
+
+## 测试结果
+|      | A100 平均生成时间（秒） | 昇腾910B 平均生成时间（秒） |
+|------|-------------------------|----------------------------|
+| 时间 | 0.0032                  | 0.0138                     |
+

dataset.json

@@ -0,0 +1 @@
+["样例数据-1", "样例数据-2", "样例数据-3"]

main.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import argparse
+import json
+import time
+from pathlib import Path
+import numpy as np
+from sentence_transformers import SentenceTransformer
+import torch
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Encode sentences with SentenceTransformer and output embeddings & pairwise cosine similarity."
+    )
+    p.add_argument("--json", help="输入文件路径（JSON，形如 ['句子1','句子2', ...]）")
+    p.add_argument("--results", help="输出文件路径（JSON）")
+    p.add_argument("--model", help="模型路径或模型名，如 BAAI/bge-large-zh-v1.5 或本地目录")
+    p.add_argument("--device", default=None,
+                   help="设备：cuda / cpu / npu；默认自动检测（优先 cuda，其次 cpu；也可显式传 npu）")
+    p.add_argument("--batch-size", type=int, default=32, help="encode 批大小，默认 32")
+    p.add_argument("--no-normalize", action="store_true", help="不做 L2 归一化（默认会归一化）")
+    args, _ = p.parse_known_args()
+    return args
+
+def auto_device(user_device: str | None) -> str:
+    if user_device:
+        if user_device == "cuda" and not torch.cuda.is_available():
+            if torch.npu.is_available():
+                return "npu"
+        return user_device
+    try:
+        if torch.cuda.is_available():
+            return "cuda"
+        if torch.npu.is_available():
+            return "npu"
+    except Exception:
+        pass
+    return "cpu"
+
+def main():
+    args = parse_args()
+    inp_path = Path(args.json)
+    out_path = Path(args.results)
+    model_path = args.model
+    device = auto_device(args.device)
+    normalize = not args.no_normalize
+
+    # 读取输入
+    with inp_path.open("r", encoding="utf-8") as f:
+        data = json.load(f)
+    if not isinstance(data, list):
+        raise ValueError("输入 JSON 必须是数组格式，如：['句子1', '句子2', ...]")
+    sentences = [str(x) for x in data]
+
+    # 加载模型
+    model = SentenceTransformer(model_path, device=device)
+
+    # 编码并计时
+    t0 = time.time()
+    embeddings = model.encode(
+        sentences,
+        batch_size=args.batch_size,
+        normalize_embeddings=normalize,
+        convert_to_numpy=True,
+        device=device
+    )
+    encode_time = time.time() - t0
+
+    # 若未归一化，则计算相似度前先做归一化（保证 similarity 为余弦相似度）
+    if not normalize:
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-12
+        embeddings = embeddings / norms
+
+    # 两两相似度（余弦）——已归一化则点积即余弦
+    similarity = embeddings @ embeddings.T
+    avg_latency = encode_time / len(sentences) if sentences else 0
+
+    # 组织输出
+    result = {
+        "model_path": model_path,
+        "device": device,
+        "count": len(sentences),
+        "dim": int(embeddings.shape[1]) if len(embeddings.shape) == 2 else None,
+        "total_elapsed_seconds": round(float(encode_time), 6),
+        "avg_latency": avg_latency,
+        "sentences": sentences,
+        "embeddings": embeddings.tolist(),          # [N, D]
+        "similarity": similarity.tolist()           # [N, N]
+    }
+
+    # 保存
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    with out_path.open("w", encoding="utf-8") as f:
+        json.dump(result, f, ensure_ascii=False, indent=2)
+
+    print(f"✅ Done. Saved to: {out_path}")
+
+if __name__ == "__main__":
+    main()
+

run_in_docker.sh

@@ -0,0 +1,4 @@
+#! /usr/bin/env bash
+image=feature:v0.1
+device=1
+docker run -v `pwd`:/workspace -e ASCEND_VISIBLE_DEVICES=$device -e NPU_VISIBLE_DEVICES=${device} --device /dev/davinci$device:/dev/davinci0 --device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc -v /mnt:/mnt -v /usr/local/dcmi:/usr/local/dcmi -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info -v /etc/ascend_install.info:/etc/ascend_install.info --privileged --entrypoint bash  $image test.sh

test.sh

@@ -0,0 +1 @@
+python main.py --json dataset.json --results output.json --model /mnt/contest_ceph/zhanghao/models/BAAI/bge-large-zh-v1.5 --device cuda