Skip to content

Commit

Permalink
fix
Browse files Browse the repository at this point in the history
  • Loading branch information
Mini256 committed Mar 7, 2025
1 parent 760b74e commit 77f95f2
Showing 1 changed file with 52 additions and 49 deletions.
101 changes: 52 additions & 49 deletions core/tests/tidb-client-quickstart.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": 39,
"id": "4995a54f311c4b1c",
"metadata": {},
"outputs": [],
Expand All @@ -100,7 +100,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 40,
"id": "bdddb9f0a005b74d",
"metadata": {},
"outputs": [],
Expand All @@ -113,11 +113,13 @@
"# Define your embedding model.\n",
"text_embed = EmbeddingFunction(\"openai/text-embedding-3-small\")\n",
"\n",
"class Chunk(TiDBModel, table=True):\n",
"class Chunk(TiDBModel, table=True): \n",
" __tablename__ = \"chunks\"\n",
" __table_args__ = {'extend_existing': True}\n",
"\n",
" id: int = Field(primary_key=True)\n",
" text: str = Field()\n",
" text_vec: Optional[Any] = text_embed.VectorField(source_field=\"text\")\n",
" text_vec: Optional[Any] = text_embed.VectorField(source_field=\"text\") # 👈 Define the vector field.\n",
" user_id: int = Field()\n",
"\n",
"table = db.create_table(schema=Chunk)"
Expand All @@ -128,27 +130,24 @@
"id": "3eab5d6eaaaaa868",
"metadata": {},
"source": [
"### Insert Data"
"### Insert Data\n",
"\n",
"🔢 Auto embedding: when you insert new data, the SDK automatically embeds the corpus for you."
]
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 46,
"id": "baec9a5ae06231be",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[Chunk(user_id=2, id=2, text='A quick brown dog runs in the park', text_vec=array([-0.0412815 , -0.00934362, 0.01239674, ..., -0.00587278,\n",
" -0.00735941, 0.01383422], dtype=float32)),\n",
" Chunk(user_id=2, id=3, text='The lazy fox sleeps under the tree', text_vec=array([-0.01610469, -0.00269681, -0.01787939, ..., -0.00041015,\n",
" 0.01320426, 0.02987844], dtype=float32)),\n",
" Chunk(user_id=3, id=4, text='A dog and a fox play in the park', text_vec=array([-2.7123539e-02, -4.4581316e-02, -3.8457386e-02, ...,\n",
" -1.1360981e-03, 9.5597192e-05, 3.4092940e-02], dtype=float32))]"
"4"
]
},
"execution_count": 5,
"execution_count": 46,
"metadata": {},
"output_type": "execute_result"
}
Expand All @@ -159,7 +158,8 @@
" Chunk(id=2, text=\"A quick brown dog runs in the park\", user_id=2),\n",
" Chunk(id=3, text=\"The lazy fox sleeps under the tree\", user_id=2),\n",
" Chunk(id=4, text=\"A dog and a fox play in the park\", user_id=3)\n",
"])"
"])\n",
"table.rows()"
]
},
{
Expand All @@ -172,29 +172,36 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": 53,
"id": "3c4313022f06bd3e",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[(4, 'A dog and a fox play in the park', 0.7308190419242949),\n",
" (2, 'A quick brown dog runs in the park', 0.665493189763966),\n",
" (1, 'The quick brown fox jumps over the lazy dog', 0.6157064668170177)]"
"[('A quick brown dog runs in the park', 0.665493189763966),\n",
" ('The lazy fox sleeps under the tree', 0.554631888866523)]"
]
},
"execution_count": 8,
"execution_count": 53,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"chunks = table.search(\"A quick fox in the park\").limit(3).to_pydantic()\n",
"[\n",
" (c.id, c.text, c.score)\n",
" for c in chunks\n",
"]"
"from autoflow.storage.tidb import DistanceMetric\n",
"\n",
"chunks = (\n",
" table.search(\"A quick fox in the park\") # 👈 The query will be embedding automatically.\n",
" # .distance_metric(metric=DistanceMetric.COSINE)\n",
" # .num_candidate(20)\n",
" .filter({\n",
" \"user_id\": 2\n",
" })\n",
" .limit(2)\n",
" .to_pydantic()\n",
")\n",
"[(c.text, c.score) for c in chunks]"
]
},
{
Expand All @@ -221,7 +228,7 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": 49,
"id": "ace02b45",
"metadata": {},
"outputs": [
Expand All @@ -231,7 +238,7 @@
"[(1, 'The quick brown fox jumps over the lazy dog', 1)]"
]
},
"execution_count": 11,
"execution_count": 49,
"metadata": {},
"output_type": "execute_result"
}
Expand All @@ -249,39 +256,35 @@
"id": "af9c3428",
"metadata": {},
"source": [
"### Truncate table"
"### Truncate table\n",
"\n",
"Clear all data in the table:"
]
},
{
"cell_type": "code",
"execution_count": 12,
"execution_count": 45,
"id": "cceb0bf0",
"metadata": {},
"outputs": [],
"outputs": [
{
"data": {
"text/plain": [
"0"
]
},
"execution_count": 45,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"table.truncate()"
"table.truncate()\n",
"table.rows()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
}
},
"metadata": {},
"nbformat": 5,
"nbformat_minor": 9
}

0 comments on commit 77f95f2

Please sign in to comment.