diff --git a/tools/train/train/qwq_infer.ipynb b/tools/train/train/qwq_infer.ipynb new file mode 100644 index 0000000..200ca37 --- /dev/null +++ b/tools/train/train/qwq_infer.ipynb @@ -0,0 +1,170 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "from transformers import AutoModelForCausalLM, AutoTokenizer" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "model_path = \"/root/models/notebook\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "model = AutoModelForCausalLM.from_pretrained(\n", + " model_path,\n", + " return_dict=True,\n", + " torch_dtype=\"auto\",\n", + " device_map=\"auto\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "tokenizer = AutoTokenizer.from_pretrained(model_path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "pipeline = pipeline(\n", + " \"text-generation\",\n", + " model=model,\n", + " tokenizer=tokenizer,\n", + " return_full_text=False,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "prompt = \"How many r's are in the word \\\"strawberry\\\"\"\n", + "messages = [\n", + " {\"role\": \"user\", \"content\": prompt}\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "text = tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=False,\n", + " add_generation_prompt=True\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "streamer = TextIteratorStreamer(\n", + " tokenizer,\n", + " skip_prompt=True,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "generation_thread = Thread(target=self._pipeline, kwargs=dict(\n", + " text_inputs=text,\n", + " do_sample=True,\n", + " temperature=0.7,\n", + " max_new_tokens=512,\n", + " streamer=streamer,\n", + "))\n", + "generation_thread.start()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [ + "\n", + "for text in streamer:\n", + " print(text)" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/tools/train/train/qwq.ipynb b/tools/train/train/qwq_train.ipynb similarity index 100% rename from tools/train/train/qwq.ipynb rename to tools/train/train/qwq_train.ipynb