(tables: list[dict], question: str)
| 11 | |
| 12 | |
| 13 | def get_table_embedding(tables: list[dict], question: str): |
| 14 | _list = [] |
| 15 | for table in tables: |
| 16 | _list.append({"id": table.get('id'), "schema_table": table.get('schema_table'), "cosine_similarity": 0.0}) |
| 17 | |
| 18 | if _list: |
| 19 | try: |
| 20 | text = [s.get('schema_table') for s in _list] |
| 21 | |
| 22 | model = EmbeddingModelCache.get_model() |
| 23 | start_time = time.time() |
| 24 | results = model.embed_documents(text) |
| 25 | end_time = time.time() |
| 26 | SQLBotLogUtil.info(str(end_time - start_time)) |
| 27 | |
| 28 | q_embedding = model.embed_query(question) |
| 29 | for index in range(len(results)): |
| 30 | item = results[index] |
| 31 | _list[index]['cosine_similarity'] = cosine_similarity(q_embedding, item) |
| 32 | |
| 33 | _list.sort(key=lambda x: x['cosine_similarity'], reverse=True) |
| 34 | _list = _list[:settings.TABLE_EMBEDDING_COUNT] |
| 35 | # print(len(_list)) |
| 36 | SQLBotLogUtil.info(json.dumps(_list)) |
| 37 | return _list |
| 38 | except Exception: |
| 39 | traceback.print_exc() |
| 40 | return _list |
| 41 | |
| 42 | |
| 43 | def calc_table_embedding(tables: list[dict], question: str): |
nothing calls this directly
no test coverage detected