在SQL Chain中使用开源LLM

Question

是否可以在SQL链中使用开源LLM模型？ 我尝试在 SQL Chain 中使用 Tapex/Flan 模型，但在 dict[] 类上出现序列化错误。

错误：

File "pydantic\main.py", line 341, in pydantic.main.BaseModel.init
pydantic.error_wrappers.ValidationError: 1 validation error for SQLDatabaseChain
root -> llm
value is not a valid dict (type=type_error.dict)

示例片段

tokenizer = TapexTokenizer.from_pretrained("microsoft/tapex-base-finetuned-wtq")
model = BartForConditionalGeneration.from_pretrained("microsoft/tapex-base-finetuned-wtq")
chain = SQLDatabaseChain(llm=model, database=db, verbose=True)
chain.run("context query ?")

是否有任何示例/片段可用于在 SQL Chain 中使用开源 LLM 模型？

Answer 1

您是否尝试过通过在模板中描述数据库和示例数据来自己创建 sql 查询？ 然后创建一个顺序链或多链来验证sql查询是否有效，然后自己使用SQL Agent或psycog2执行查询？

Answer 2

    db = SQLDatabase.from_uri("postgresql://username:password@localhost:5432/db_name?connect_timeout=10&target_session_attrs=primary",
                                  include_tables=['table_name'])
    chain = SQLDatabaseChain(llm=CustomLLM(), database=db, verbose=True)
    chain.run("How many users are there?")
        
    class CustomLLM(LLM):
        model_name = "microsoft/tapex-base-finetuned-wtq"
        tokenizer = TapexTokenizer.from_pretrained("microsoft/tapex-base-finetuned-wtq")
        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        model = BartForConditionalGeneration.from_pretrained("microsoft/tapex-base-finetuned-wtq")
    
        def _call(self, prompt, stop=None, **kwargs):
            prompt = f"### User:\n{prompt}\n\n### Assistant:\n"
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
            output = self.model.generate(**inputs, streamer=self.streamer, use_cache=True, max_new_tokens=float('inf'))
            return self.tokenizer.decode(output[0], skip_special_tokens=True)
    
        @property
        def _identifying_params(self):
            return {"name_of_model": self.model_name}
    
        @property
        def _llm_type(self):
            return "custom"