Newest Embedding model

[chengzicong20040913]

class EmbeddingModel:
    """Base class for Embedding models"""
    def __init__(self, model: str,openai_api_key: str = None, openai_api_base: str = None):
        self.model_name = model
        self.key = openai_api_key
        self.api_base = openai_api_base
        self.client = None
        if self.model_name == 'OpenAI':
            self.client = OpenAIEmbeddings(
                model=self.model_name,
                openai_api_key=self.key,
                openai_api_base=self.api_base
            )
        else:
            self.client = OpenAI(
                api_key=self.key,
                base_url=self.api_base,
            )

    def embed_query(self, text: str):
        if self.model_name == 'OpenAI':
            return self.client.embed_query(text)
        else:
            completion = self.client.embeddings.create(
                model=self.model_name,
                input=[text],
                dimensions=1024,
                encoding_format="float"
            )
            output=completion.model_dump_json()
            output = json.loads(output)
            # 1. 先按index排序
            sorted_data = sorted(output["data"], key=lambda x: int(x["index"]))
            embeddings = [item["embedding"] for item in sorted_data]
            return embeddings[0]
            

    def embed_documents(self, texts: List[str]) -> List[np.ndarray]:
        for text in texts:
            if self.model_name == 'OpenAI':
                return self.client.embed_documents(texts)
            else:
                #按照10,10分片的方式进行分割
                for i in range(0, len(texts), 10):
                    chunk = texts[i:i + 10]
                    #对于chunk需要对于每一个条目裁剪成最大8192个字符
                    chunk = [text[:16384] for text in chunk]
                    # 2. 分片处理
                    completion = self.client.embeddings.create(
                        model=self.model_name,
                        input=chunk,
                        dimensions=1024,
                        encoding_format="float"
                    )
                    output=completion.model_dump_json()
                    output = json.loads(output)
                    # 1. 先按index排序
                    sorted_data = sorted(output["data"], key=lambda x: int(x["index"]))
                    embeddings = [item["embedding"] for item in sorted_data]
                    # 3. 拼接
                    if i == 0:
                        all_embeddings = embeddings
                    else:
                        all_embeddings.extend(embeddings)
                return all_embeddings
            
    async def aembed_query(self, text: str) -> List[float]:
        """异步生成单个文本的embedding"""
        if self.model_name == 'OpenAI':
            # 假设OpenAIEmbeddings有异步方法
            return await self.client.aembed_query(text)
        else:
            async with aiohttp.ClientSession() as session:
                headers = {
                    "Authorization": f"Bearer {self.key}",
                    "Content-Type": "application/json"
                }
                payload = {
                    "model": self.model_name,
                    "input": [text],
                    "dimensions": 1024,
                    "encoding_format": "float"
                }
                
                async with session.post(
                    f"{self.api_base}/embeddings",
                    headers=headers,
                    json=payload
                ) as response:
                    output = await response.json()
                    sorted_data = sorted(output["data"], key=lambda x: int(x["index"]))
                    return [item["embedding"] for item in sorted_data][0]

    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        """异步批量生成embedding"""
        if self.model_name == 'OpenAI':
            return await self.client.aembed_documents(texts)
        else:
            # 使用asyncio.gather并发处理
            tasks = [self.aembed_query(text) for text in texts]
            return await asyncio.gather(*tasks)

[code4DB]

Please submit a pull request to facilitate the understanding of the code changes :)

[ZZH-qwq]

@code4DB 您好！我在进行 CrackSQL 的复现实验时也遇到了相似的问题。因此提供更多信息作为补充：

chengzicong 同学提到的问题原因来自 调用来自阿里云的 Embedding model (例如 text-embedding-v3 ) 时，无法通过langchain_openai.OpenAIEmbeddings 实现 aembed_documents ，会提示参数错误：

{'code': 'InvalidParameter', 'param': None, 'message': '<400> InternalError.Algo.InvalidParameter: Value error, contents is neither str nor list of str.: input.contents', 'type': 'InvalidParameter'}

解决方法参考（在 错误查询页面 中查找上述信息得到） https://help.aliyun.com/zh/model-studio/use-bailian-in-langchain ，改为使用 DashScopeEmbeddings 作为 embedding_model。这样无需自定义 EmbeddingModel 类即可（一定程度上）解决问题：

(如图，成功获得了返回值)

我的具体修改为：在 backend.llm_model.embeddings #L89-L105：

CrackSQL/backend/llm_model/embeddings.py

Lines 89 to 105 in 0b2112f

	if model_config['deployment_type'] == 'cloud':
	# Cloud model (OpenAI)
	embedding = OpenAIEmbeddings(
	model=model_config['name'],
	openai_api_key=model_config['api_key'],
	openai_api_base=model_config['api_base']
	)
	else:
	import torch
	# Determine device
	device = "cuda" if torch.cuda.is_available() else "cpu"
	# Local model (HuggingFace)
	embedding = HuggingFaceEmbeddings(
	model_name=model_config['model_path'],
	model_kwargs={'device': device},
	encode_kwargs={'normalize_embeddings': True}
	)

添加针对阿里云模型的初始化方式（同时修改了 init_config.yaml 的相关设置）：

elif model_config["deployment_type"] == "cloud-aliyun":
    # Cloud model (Aliyun)
    embedding = DashScopeEmbeddings(
        model=model_config["name"],
        dashscope_api_key=model_config["api_key"],
    )

但是目前仍然遇到了以下问题：

Failed to generate Embedding: status_code: 400 
 code: InvalidParameter
 message: Range of input length should be [1, 8192]

这是由于被向量化的单条文本 token 长度超过了限制（原代码使用的 text-embedding-ada-002 上限为 8191，我使用的 text-embedding-v3 上限为 8192）：

我注意到 @chengzicong20040913 同学在代码中使用了截断进行 token 数量控制：

def embed_documents(self, texts: List[str]) -> List[np.ndarray]:
    ...
                #对于chunk需要对于每一个条目裁剪成最大8192个字符
                chunk = [text[:16384] for text in chunk]

但这不太合理：上图第二次请求的内容最大字符数为 22565，但 token 总数在限制之内。因此这样截取可能损失了部分有效信息。

以下是我记录的超过限制的文本（共 35212 字符，文件的完整内容即为传入 aembed_documents 的 List[str] 中的一项）：

max_len_text.txt

我使用 OpenAI token 计数器 统计了它的 token 数量，发现它也超过了代码中使用过的 text-embedding-ada-002 模型的限制：

因此我希望询问开发团队以下信息：

• 请问在进行测试的过程中使用的是什么文本嵌入模型，是否有遇到类似的问题？
• 这个问题是否在换用其它（例如本地 / OpenAI）模型后能够解决？（因为我注意到其它 issue 中没有提到达到 token 限制问题）
• 如果仍然使用阿里云的词嵌入模型，解决方法有比截断文本更好的选择吗？

希望能够得到 @code4DB 或是其它开发人员的回复。如果需要，我很乐意提供更多信息。感谢您为项目的付出。

[code4DB]

@ZZH-qwq 感谢您的详细回复！

请问在进行测试的过程中使用的是什么文本嵌入模型，是否有遇到类似的问题？

当前测试支持的模型为Huggingface上的主流开源embedding模型，以及满足OpenAI接口的embedding模型，其余模型需要做适配包括调用的函数以及接口（正如提到的阿里云的embedding模型）。可以将您针对模型适配的修改提交一个pull request，经验证后我们会合入主分支。

这个问题是否在换用其它（例如本地 / OpenAI）模型后能够解决？（因为我注意到其它 issue 中没有提到达到 token 限制问题）

该问题源于输入给模型的token数过大，因此最直接的方法为换用支持超长token数的embedding模型来得到解决

如果仍然使用阿里云的词嵌入模型，解决方法有比截断文本更好的选择吗？

正如你所提供的文档，token数的主要占比源于检索到的文本信息，这些文本信息的篇幅首先由超参topk控制（即意味着应选择适中的topk值，例如2-3）。此外，在我们原来的方法中会对检索到的文档的不同部分进行截断处理（即只保留核心的部分）。需要再次强调的是，该问题的核心关乎到检索到的文本信息的质量问题（例如，删减繁冗的描述）。另外，注意到提供的附件文件包含很多重复的文档。清理以及添加这些语法文档在我们的未来规划中，也欢迎大家提供高质量的文档：）

[ZZH-qwq]

@code4DB 感谢您的回复~

可以将您针对模型适配的修改提交一个pull request，经验证后我们会合入主分支。

好的。我会在进一步测试之后提交 PR。但目前添加的内容涉及引入新的依赖项（dashscope），不过在不实例化 DashScopeEmbeddings 的情况下无需安装。因此仅作说明，我不会将它写入 requirements.txt。

换用支持超长token数的embedding模型来得到解决

我了解到目前大部分 Embedding Model 的最长 token 都是 8192 左右，其它模型能够正常运行是因为 token 计算方式不同吗？

另外，注意到提供的附件文件包含很多重复的文档。

非常抱歉我上传了错误的文档（它记录了三条等长但内容相似的查询），可能造成了部分误解。正确的请求内容如下：

max_len_text.txt

[code4DB]

@ZZH-qwq Could you confirm whether the issue here occurred during the initialization of the knowledge base? Regarding the initialization of the knowledge base, we did not encounter any token limit-related problems during testing. Additionally, splitting long texts into chunks is currently a common processing method.

[ZZH-qwq]

@code4DB 是的，这个问题确实出现在初始化 Knowledge base 的时候。

python init_knowledge_base.py  --config_file .\config\init_config.yaml

我怀疑可能是不同模型 token 划分有些区别。因此这个问题可能是来自于阿里云的词嵌入模型。我在将传入的字符串截断至 24000 个字符之后没有再遇到其它问题，成功完成了初始化过程。

---

Yes, this issue does occur during the initialization of the Knowledge base.

python init_knowledge_base.py  --config_file .\config\init_config.yaml

I suspect it might be due to some differences in tokenization between different models. Therefore, this problem likely originates from Aliyun's text embedding model. After truncating the input string to 24,000 characters, I haven't encountered any further issues and successfully completed the initialization process.

[code4DB]

@ZZH-qwq It seems weird. You can refer to the files utilized to initialize the knowledge base in data/processed_document. Apparently, I think not a single item has the risk to exceed the token number.

[ZZH-qwq]

@code4DB 之前回复中所附带的文件确实是由提供的知识库生成的。它来自于：

data/processed_document/mysql_8_kb.json

的第 3998 - 4020 行。源文件无法直接在此处展示完全，仅展示每项的字符数：

{
    "type": "type",
    "keyword": "TIMESTAMP ( fsp )",
    "tree": "(dataType TIMESTAMP (lengthOneDimension ( )))",
    "description": [
        "MySQL permits fractional seconds for <link><code>TIME ...(1201 characters)",
        "Any <link><code>TIMESTAMP</code>&LINK&https://dev ...(458 characters)",
        " A timestamp. The range is <code>'1970-01-01 ...(3754 characters)"
    ],
    "detail": "[DESCRIPTION]: MySQL permits fractional seconds for ...(29997 characters)",
    "example": [],
    "link": [
        "https://dev.mysql.com/doc/refman/8.0/en/date-and-time-types.html",
        "...(8 more links)"
    ]
},

注意到我上传的文件完整包含了其中的 "detail"、"description" 项（仅部分字符被转义），因此总字符数远大于 30000 个。

附上 diff 对比节选（调换了两者在 json 中的顺序）：

• "detail" 的起始部分：

  

• "detail" 的结尾，"description" 的起始部分：

  

[code4DB]

@ZZH-qwq Hello, the model is restricted by the number of tokens but not the characters. We will check such issues.

[ZZH-qwq]

the model is restricted by the number of tokens but not the characters.

@code4DB That's correct. However in my previous test, the input still hit the maximum token limit. (Since no token counter is found for Aliyun's embedding models, I used OpenAI's for reference):

我使用 OpenAI token 计数器 统计了它的 token 数量，发现它也超过了代码中使用过的 text-embedding-ada-002 模型的限制：

[HuanleTT]

@ZZH-qwq 同学你好，我这边在修改init_config.yaml文件后还是出现相同的报错，能请问你的init_config.yaml文件是怎么配置的吗？谢谢！

[ZZH-qwq]

@HuanleTT 参考 #15 中提到的问题，程序会读取之前已经存在的配置，因此修改的 cloud-aliyun 等参数不会生效。

我个人建议删除 backend/instance 目录 (其中包含存储配置的 info.db)，然后重新运行 init_knowledge_base.py。

我的 init_config.yaml 配置如下：

EMBEDDING_MODELS:
  - name: "text-embedding-v3"
    deployment_type: "cloud-aliyun"
    category: "embedding"
    api_base: "https://dashscope.aliyuncs.com/compatible-mode/v1"
    api_key: "**************"
    path: ""
    dimension: 1024
    description: "Embedding模型，用于生成文本向量"

[HuanleTT]

@ZZH-qwq 根据你的回复以及参考官方文档，我成功修改了配置文件，在token数量限制内导入了部分知识库。非常感谢！

[SJTUer-sxr]

@ZZH-qwq 同学你好，我在使用"text-embedding-v3"模型时遇到了同样的问题，通过将传入的字符串截断至24000确实可以避免在初始化知识库时发生错误，但在web端运行对sql语句进行翻译时依然会出现这种token长度超过限制的问题，请问你是否遇到了这种问题，有没有好一点的解决方法。

[chengzicong20040913]

我觉得应该多截断一点到16384应该不会报错

[ZZH-qwq]

@SJTUer-sxr 这个报错并非来自 Embedding model，而是来自 LLM model。

建议检查您的配置文件 init_config.yaml > LLM_MODELS > (model) > max_tokens 项目（通常位于 LLM_MODELS 配置的最后一行：

CrackSQL/backend/config/init_config.yaml.copy

Line 14 in 56cf0ed

max_tokens: 128000 # Maximum tokens for generation

如果您填写的 max_tokens 数量大于模型支持的最大回复长度（根据报错应当是 8192）那么就会出现这个问题。需要修改成 [1, 8192] 之间的值。

同时这个问题也有可能来自此前初始化时参数被写入数据库，然后修改配置文件不会覆盖已初始化的内容，导致错误的 max_tokens 仍然生效。

因此建议按照上述步骤检查和修改 init_config.yaml 之后，删除 backend/instance 目录，然后重新运行 init_knowledge_base.py（这需要重新向量化知识库）。问题应当能够被解决。

[SJTUer-sxr]

@ZZH-qwq @chengzicong20040913 根据你们的回复，我成功解决了这个问题，万分感谢！