Spaces:
Paused
Paused
File size: 4,865 Bytes
3161850 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """
embedding_client.py — 语义向量召回客户端
支持两种模式：
1. API 模式（DeepSeek / OpenAI / Qwen 兼容嵌入接口）
2. 本地 sentence-transformers 模式（可选，需安装）
优先级：API → sentence-transformers → token cosine fallback
"""
from __future__ import annotations
import json
import os
import re
from pathlib import Path
from typing import Optional
try:
from dotenv import load_dotenv
except ModuleNotFoundError: # python-dotenv is optional; env vars still work.
def load_dotenv(*args, **kwargs):
return False
load_dotenv()
class EmbeddingClient:
"""统一嵌入接口。未配置时自动回退。"""
def __init__(self) -> None:
self.api_key = os.getenv("EMBEDDING_API_KEY") or os.getenv("LLM_API_KEY")
self.base_url = os.getenv("EMBEDDING_BASE_URL") or os.getenv("LLM_BASE_URL")
self.model = os.getenv("EMBEDDING_MODEL", "text-embedding-v1")
self.provider = os.getenv("EMBEDDING_PROVIDER", "").lower()
self._client = None
self._local_model = None
self.mode: Optional[str] = None # "api" | "local" | None
# 尝试加载 API 模式
if self.api_key and self.base_url:
try:
from openai import OpenAI
self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)
self.mode = "api"
print(f"[OK] EmbeddingClient: API 模式 ({self.base_url})")
except Exception as exc:
print(f"[WARN] EmbeddingClient API 初始化失败：{exc}")
# API 失败则尝试本地模型
if self.mode is None:
self._try_load_local()
def _try_load_local(self) -> None:
try:
from sentence_transformers import SentenceTransformer
model_name = os.getenv("LOCAL_EMBEDDING_MODEL", "BAAI/bge-small-zh-v1.5")
self._local_model = SentenceTransformer(model_name)
self.mode = "local"
print(f"[OK] EmbeddingClient: 本地模式 ({model_name})")
except Exception:
self.mode = None
def is_available(self) -> bool:
return self.mode is not None
def embed(self, texts: list[str]) -> Optional[list[list[float]]]:
"""批量获取嵌入向量。返回 List[List[float]] 或 None。"""
if self.mode == "api":
return self._embed_api(texts)
if self.mode == "local":
return self._embed_local(texts)
return None
def _embed_api(self, texts: list[str]) -> Optional[list[list[float]]]:
try:
resp = self._client.embeddings.create(model=self.model, input=texts)
return [d.embedding for d in resp.data]
except Exception as exc:
print(f"[WARN] Embedding API 调用失败：{exc}")
return None
def _embed_local(self, texts: list[str]) -> Optional[list[list[float]]]:
try:
embeddings = self._local_model.encode(texts, normalize_embeddings=True)
return embeddings.tolist()
except Exception as exc:
print(f"[WARN] 本地嵌入失败：{exc}")
return None
def cosine_similarity(self, vec1: list[float], vec2: list[float]) -> float:
"""余弦相似度。"""
if not vec1 or not vec2:
return 0.0
dot = sum(a * b for a, b in zip(vec1, vec2))
norm1 = sum(a * a for a in vec1) ** 0.5
norm2 = sum(b * b for b in vec2) ** 0.5
if norm1 == 0 or norm2 == 0:
return 0.0
return dot / (norm1 * norm2)
def semantic_similarity(self, resume_text: str, job_text: str) -> Optional[float]:
"""计算简历与岗位的语义相似度（0~1）。"""
embeddings = self.embed([resume_text, job_text])
if embeddings is None:
return None
return self.cosine_similarity(embeddings[0], embeddings[1])
def rank_by_semantic(
self, resume_text: str, jobs: list[dict], top_k: int = 8
) -> list[tuple[int, float]]:
"""用语义向量对岗位重新排序。返回 [(index, score), ...] 按分数降序。"""
if not self.is_available():
return []
texts = [f"{j.get('title', '')} {j.get('jd', '')}" for j in jobs]
all_texts = [resume_text] + texts
embeddings = self.embed(all_texts)
if embeddings is None:
return []
resume_emb = embeddings[0]
scores: list[tuple[int, float]] = []
for i, job_emb in enumerate(embeddings[1:], start=0):
sim = self.cosine_similarity(resume_emb, job_emb)
scores.append((i, sim))
scores.sort(key=lambda x: x[1], reverse=True)
return scores[:top_k]
def get_embedding_client() -> EmbeddingClient:
"""工厂函数。"""
return EmbeddingClient()
|