add sample notebook preload
This commit is contained in:
parent
8ebe3e6b5b
commit
c27000c787
|
|
@ -1,4 +1,5 @@
|
||||||
transformers
|
transformers
|
||||||
torch
|
torch
|
||||||
huggingface_hub
|
huggingface_hub
|
||||||
pathlib
|
numpy
|
||||||
|
ipywidgets
|
||||||
|
|
@ -0,0 +1,104 @@
|
||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 5,
|
||||||
|
"id": "8048aa56-4549-4afa-b8b0-d111cc7020c3",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"[[0.7645573019981384, 0.14142529666423798], [0.13549786806106567, 0.5999557375907898]]\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"# Requires transformers>=4.51.0\n",
|
||||||
|
"\n",
|
||||||
|
"import torch\n",
|
||||||
|
"import torch.nn.functional as F\n",
|
||||||
|
"\n",
|
||||||
|
"from torch import Tensor\n",
|
||||||
|
"from transformers import AutoTokenizer, AutoModel\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def last_token_pool(last_hidden_states: Tensor,\n",
|
||||||
|
" attention_mask: Tensor) -> Tensor:\n",
|
||||||
|
" left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])\n",
|
||||||
|
" if left_padding:\n",
|
||||||
|
" return last_hidden_states[:, -1]\n",
|
||||||
|
" else:\n",
|
||||||
|
" sequence_lengths = attention_mask.sum(dim=1) - 1\n",
|
||||||
|
" batch_size = last_hidden_states.shape[0]\n",
|
||||||
|
" return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def get_detailed_instruct(task_description: str, query: str) -> str:\n",
|
||||||
|
" return f'Instruct: {task_description}\\nQuery:{query}'\n",
|
||||||
|
"\n",
|
||||||
|
"# Each query must come with a one-sentence instruction that describes the task\n",
|
||||||
|
"task = 'Given a web search query, retrieve relevant passages that answer the query'\n",
|
||||||
|
"\n",
|
||||||
|
"queries = [\n",
|
||||||
|
" get_detailed_instruct(task, 'What is the capital of China?'),\n",
|
||||||
|
" get_detailed_instruct(task, 'Explain gravity')\n",
|
||||||
|
"]\n",
|
||||||
|
"# No need to add instruction for retrieval documents\n",
|
||||||
|
"documents = [\n",
|
||||||
|
" \"The capital of China is Beijing.\",\n",
|
||||||
|
" \"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.\"\n",
|
||||||
|
"]\n",
|
||||||
|
"input_texts = queries + documents\n",
|
||||||
|
"\n",
|
||||||
|
"tokenizer = AutoTokenizer.from_pretrained('/tmp/models-cache/Qwen/Qwen3-Embedding-0.6B', padding_side='left')\n",
|
||||||
|
"model = AutoModel.from_pretrained('/tmp/models-cache/Qwen/Qwen3-Embedding-0.6B')\n",
|
||||||
|
"\n",
|
||||||
|
"# We recommend enabling flash_attention_2 for better acceleration and memory saving.\n",
|
||||||
|
"# model = AutoModel.from_pretrained('Qwen/Qwen3-Embedding-0.6B', attn_implementation=\"flash_attention_2\", torch_dtype=torch.float16).cuda()\n",
|
||||||
|
"\n",
|
||||||
|
"max_length = 8192\n",
|
||||||
|
"\n",
|
||||||
|
"# Tokenize the input texts\n",
|
||||||
|
"batch_dict = tokenizer(\n",
|
||||||
|
" input_texts,\n",
|
||||||
|
" padding=True,\n",
|
||||||
|
" truncation=True,\n",
|
||||||
|
" max_length=max_length,\n",
|
||||||
|
" return_tensors=\"pt\",\n",
|
||||||
|
")\n",
|
||||||
|
"batch_dict.to(model.device)\n",
|
||||||
|
"outputs = model(**batch_dict)\n",
|
||||||
|
"embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])\n",
|
||||||
|
"\n",
|
||||||
|
"# normalize embeddings\n",
|
||||||
|
"embeddings = F.normalize(embeddings, p=2, dim=1)\n",
|
||||||
|
"scores = (embeddings[:2] @ embeddings[2:].T)\n",
|
||||||
|
"print(scores.tolist())\n",
|
||||||
|
"# [[0.7645568251609802, 0.14142508804798126], [0.13549736142158508, 0.5999549627304077]]\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "Python 3 (ipykernel)",
|
||||||
|
"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.12.10"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 5
|
||||||
|
}
|
||||||
|
|
@ -26,4 +26,16 @@ metadata:
|
||||||
helm.sh/chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
|
helm.sh/chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
|
||||||
data:
|
data:
|
||||||
download_models.py: |-
|
download_models.py: |-
|
||||||
{{ .Files.Get "files/download_models.py" | nindent 4 }}
|
{{ .Files.Get "files/download_models.py" | nindent 4 }}
|
||||||
|
---
|
||||||
|
apiVersion: v1
|
||||||
|
kind: ConfigMap
|
||||||
|
metadata:
|
||||||
|
name: {{ .Release.Name }}-welcome-notebook
|
||||||
|
labels:
|
||||||
|
app.kubernetes.io/managed-by: {{ .Release.Service | quote }}
|
||||||
|
app.kubernetes.io/instance: {{ .Release.Name | quote }}
|
||||||
|
helm.sh/chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
|
||||||
|
data:
|
||||||
|
welcome.ipynb: |-
|
||||||
|
{{ .Files.Get "files/welcome.ipynb" | nindent 4 }}
|
||||||
|
|
@ -8,6 +8,11 @@ jupyterhub:
|
||||||
|
|
||||||
singleuser:
|
singleuser:
|
||||||
fsGid: 100
|
fsGid: 100
|
||||||
|
defaultUrl: "/lab/tree/welcome.ipynb"
|
||||||
|
lifecycleHooks:
|
||||||
|
postStart:
|
||||||
|
exec:
|
||||||
|
command: ["/bin/sh", "-c", "pip install -r /tmp/requirements.txt"]
|
||||||
initContainers:
|
initContainers:
|
||||||
- name: model-initializer
|
- name: model-initializer
|
||||||
image: python:3.9-slim-bullseye
|
image: python:3.9-slim-bullseye
|
||||||
|
|
@ -47,6 +52,9 @@ jupyterhub:
|
||||||
- name: hf-download-script
|
- name: hf-download-script
|
||||||
configMap:
|
configMap:
|
||||||
name: "{{ .Release.Name }}-hf-download-script"
|
name: "{{ .Release.Name }}-hf-download-script"
|
||||||
|
- name: welcome-notebook
|
||||||
|
configMap:
|
||||||
|
name: "{{ .Release.Name }}-welcome-notebook"
|
||||||
- name: hf-token-secret
|
- name: hf-token-secret
|
||||||
secret:
|
secret:
|
||||||
secretName: "{{ .Release.Name }}-hf-token-secret"
|
secretName: "{{ .Release.Name }}-hf-token-secret"
|
||||||
|
|
@ -60,6 +68,9 @@ jupyterhub:
|
||||||
- name: hf-download-script
|
- name: hf-download-script
|
||||||
mountPath: /tmp/download_models.py
|
mountPath: /tmp/download_models.py
|
||||||
subPath: download_models.py
|
subPath: download_models.py
|
||||||
|
- name: welcome-notebook
|
||||||
|
mountPath: /tmp/welcome.ipynb
|
||||||
|
subPath: welcome.ipynb
|
||||||
- name: hf-token-secret
|
- name: hf-token-secret
|
||||||
mountPath: "/etc/secrets/huggingface"
|
mountPath: "/etc/secrets/huggingface"
|
||||||
readOnly: true
|
readOnly: true
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue