Skip to content

Instantly share code, notes, and snippets.

@Norod

Norod/hebrew-poem-generator-flaxgptneoforcausallm-inference-tpu.ipynb

Last active January 26, 2025 12:26
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save Norod/191cd987da004101dc1a07e4f9d6de95 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save Norod/191cd987da004101dc1a07e4f9d6de95 to your computer and use it in GitHub Desktop.
Download ZIP
hebrew-poem-generator-flaxgptneoforcausallm-inference-tpu.ipynb
Display the source blob
Display the rendered blob
Raw
hebrew-poem-generator-flaxgptneoforcausallm-inference-tpu.ipynb
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": [],
"gpuType": "V28",
"authorship_tag": "ABX9TyNaZ9dz/Dpf5vzzoaddL0kA",
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
},
"accelerator": "TPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/Norod/191cd987da004101dc1a07e4f9d6de95/hebrew-poem-generator-flaxgptneoforcausallm-inference-tpu.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "code",
"source": [
"#@title Set up JAX\n",
"#@markdown If you see an error, make sure you are using a TPU backend. Select `Runtime` in the menu above, then select the option \"Change runtime type\" and then select `TPU` under the `Hardware accelerator` setting.\n",
"import jax.tools.colab_tpu\n",
"!pip install flax transformers ftfy tokenizers\n",
"jax.devices()"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "QNhjBQpJwUnI",
"outputId": "d5f1fb6c-58db-4ee0-e01a-1c2bbb284387"
},
"execution_count": 1,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Requirement already satisfied: flax in /usr/local/lib/python3.11/dist-packages (0.10.2)\n",
"Requirement already satisfied: transformers in /usr/local/lib/python3.11/dist-packages (4.47.1)\n",
"Requirement already satisfied: ftfy in /usr/local/lib/python3.11/dist-packages (6.3.1)\n",
"Requirement already satisfied: tokenizers in /usr/local/lib/python3.11/dist-packages (0.21.0)\n",
"Requirement already satisfied: jax>=0.4.27 in /usr/local/lib/python3.11/dist-packages (from flax) (0.4.33)\n",
"Requirement already satisfied: msgpack in /usr/local/lib/python3.11/dist-packages (from flax) (1.1.0)\n",
"Requirement already satisfied: optax in /usr/local/lib/python3.11/dist-packages (from flax) (0.1.9)\n",
"Requirement already satisfied: orbax-checkpoint in /usr/local/lib/python3.11/dist-packages (from flax) (0.6.4)\n",
"Requirement already satisfied: tensorstore in /usr/local/lib/python3.11/dist-packages (from flax) (0.1.71)\n",
"Requirement already satisfied: rich>=11.1 in /usr/local/lib/python3.11/dist-packages (from flax) (13.9.4)\n",
"Requirement already satisfied: typing-extensions>=4.2 in /usr/local/lib/python3.11/dist-packages (from flax) (4.12.2)\n",
"Requirement already satisfied: PyYAML>=5.4.1 in /usr/local/lib/python3.11/dist-packages (from flax) (6.0.2)\n",
"Requirement already satisfied: numpy>=1.23.2 in /usr/local/lib/python3.11/dist-packages (from flax) (1.26.4)\n",
"Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from transformers) (3.17.0)\n",
"Requirement already satisfied: huggingface-hub<1.0,>=0.24.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.27.1)\n",
"Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (24.2)\n",
"Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from transformers) (2024.11.6)\n",
"Requirement already satisfied: requests in /usr/local/lib/python3.11/dist-packages (from transformers) (2.32.3)\n",
"Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.5.2)\n",
"Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.11/dist-packages (from transformers) (4.67.1)\n",
"Requirement already satisfied: wcwidth in /usr/local/lib/python3.11/dist-packages (from ftfy) (0.2.13)\n",
"Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub<1.0,>=0.24.0->transformers) (2024.12.0)\n",
"Requirement already satisfied: jaxlib<=0.4.33,>=0.4.33 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.27->flax) (0.4.33)\n",
"Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.27->flax) (0.4.1)\n",
"Requirement already satisfied: opt-einsum in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.27->flax) (3.4.0)\n",
"Requirement already satisfied: scipy>=1.10 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.27->flax) (1.13.1)\n",
"Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.11/dist-packages (from rich>=11.1->flax) (3.0.0)\n",
"Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.11/dist-packages (from rich>=11.1->flax) (2.19.1)\n",
"Requirement already satisfied: absl-py>=0.7.1 in /usr/local/lib/python3.11/dist-packages (from optax->flax) (1.4.0)\n",
"Requirement already satisfied: chex>=0.1.7 in /usr/local/lib/python3.11/dist-packages (from optax->flax) (0.1.88)\n",
"Requirement already satisfied: etils[epath,epy] in /usr/local/lib/python3.11/dist-packages (from orbax-checkpoint->flax) (1.11.0)\n",
"Requirement already satisfied: nest_asyncio in /usr/local/lib/python3.11/dist-packages (from orbax-checkpoint->flax) (1.6.0)\n",
"Requirement already satisfied: protobuf in /usr/local/lib/python3.11/dist-packages (from orbax-checkpoint->flax) (4.25.5)\n",
"Requirement already satisfied: humanize in /usr/local/lib/python3.11/dist-packages (from orbax-checkpoint->flax) (4.11.0)\n",
"Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (3.4.1)\n",
"Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (3.10)\n",
"Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (2.3.0)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (2024.12.14)\n",
"Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.11/dist-packages (from chex>=0.1.7->optax->flax) (1.0.0)\n",
"Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.11/dist-packages (from markdown-it-py>=2.2.0->rich>=11.1->flax) (0.1.2)\n",
"Requirement already satisfied: importlib_resources in /usr/local/lib/python3.11/dist-packages (from etils[epath,epy]->orbax-checkpoint->flax) (6.5.2)\n",
"Requirement already satisfied: zipp in /usr/local/lib/python3.11/dist-packages (from etils[epath,epy]->orbax-checkpoint->flax) (3.21.0)\n"
]
},
{
"output_type": "execute_result",
"data": {
"text/plain": [
"[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
" TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n",
" TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
" TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n",
" TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
" TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n",
" TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
" TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"
]
},
"metadata": {},
"execution_count": 1
}
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "YcNWd7AUvfm4",
"outputId": "d14e3ba6-2e08-498f-8422-845523703b11"
},
"outputs": [
{
"output_type": "stream",
"name": "stderr",
"text": [
"/usr/local/lib/python3.11/dist-packages/torch_xla/__init__.py:253: UserWarning: `tensorflow` can conflict with `torch-xla`. Prefer `tensorflow-cpu` when using PyTorch/XLA. To silence this warning, `pip uninstall -y tensorflow && pip install tensorflow-cpu`. If you are in a notebook environment such as Colab or Kaggle, restart your notebook runtime afterwards.\n",
" warnings.warn(\n"
]
}
],
"source": [
"import jax\n",
"from transformers import FlaxGPTNeoForCausalLM, AutoTokenizer\n",
"\n",
"model_name = 'Norod78/hebrew_poetry-gpt_neo-small'\n",
"\n",
"model = FlaxGPTNeoForCausalLM.from_pretrained(model_name)\n",
"tokenizer = AutoTokenizer.from_pretrained(model_name)\n"
]
},
{
"cell_type": "code",
"source": [
"prompt_text = \"\\n\\n\\n\"\n",
"generated_max_length = 512\n",
"seed = 42\n",
"model.seed = seed\n",
"inputs = tokenizer(prompt_text, return_tensors=\"jax\")\n",
"\n",
"#Changing temperature, top_k and top_p does not seem to change the outcome\n",
"outputs = model.generate(\n",
" input_ids = inputs[\"input_ids\"],\n",
" max_length=generated_max_length,\n",
" pad_token_id = model.config.eos_token_id,\n",
" prng_key=jax.random.PRNGKey(seed),\n",
" top_k=50,\n",
" top_p=1.0,\n",
" do_sample=True,\n",
" no_repeat_ngram_size=4)\n",
"\n",
"output_sequence = outputs['sequences'].squeeze(0)\n",
"text = tokenizer.decode(output_sequence, clean_up_tokenization_spaces=True)\n",
"\n",
"print(text)\n"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "tKyc4aPz481d",
"outputId": "55cfd1e8-001c-4eec-8bde-23f1cbc4b13b"
},
"execution_count": 3,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"\n",
"\n",
"\n",
"שם היצירה: מלים מתות\n",
"מאת: אור גת\n",
"\n",
"\"לא\" זה אני\n",
"אני רואה את המלים מתות ולא מתות, אני מרגיש שיום אחד מישהו ילך\n",
"\"את מתאבדת?\" הם אומרים לי.\n",
"\"מה אתה רוצה שאעשה?\" אני עונה להם, מחמיר צעד... \"את\n",
"מלבדי? מאיפה זה בא פתאום.\n",
"\"אנחנו לא צריכות יותר את הפוזה שלך\" הם אומרים לי  ומתפרדים\n",
"\"אנחנו חייבים לצאת מהחיים האלה\" אני אומר בקול חנוק.\n",
"זה לא אני. אני הולך לדוגמא.\n",
"אני יכול להרגיש שוב רק את עצמי ולא את הגוף, ולא את הנפש.\n",
"\"לא\", אני אומר להם אבל אני לא שומע אותם.\n",
"\"שום דבר לא השתנה\" הם אומרים בקול רועד.\n",
"הם יבהירו לי מה קרה, מה בדיוק קרה.\n",
"הדלתות.\n",
"\n",
"\n",
"\n",
"-שם היצירה: שיר ליום ההולדת\n",
"מאת: היצירה האלפית היחידה\n",
"\n",
"לא יודעת מתי ישוב לי פתאום. אני מתאבדת\n",
"אך אני כאן לבד\n",
"לא יראה אותי מחר. אני לא מכירה את החיים\n",
"אני תוהה אם הוא יברח\n",
"הוא קורא בקול שאני לא שומעת\n",
"הוא שוב מסתובב. לא יודעת לאן לאן לפנות\n",
"והזמן כמו זוכר. מחר הוא יבוא\n",
"אני לא מבינה איך אני הולכת לו. מחר יהיה טוב...\n",
"אני שוכבת במיטה, ערה.\n",
"כל פעם שהוא קורא דבר אני לא מפסיקה לבכות\n",
"ואני מתהאבדת.\n",
"ואני מדממת כרגע\n",
"לא רוצה לקום.\n",
"אך אני שוכבת בשקט.\n",
"מנסה לאחוז בשמיכה\n",
"הוא מפוחד כל כך שאני עומדת ככה\n",
"הוא יושב עם סיגריה בפה שלו ואני מדממת לו את הלב.\n",
"תמיד חשבתי שהוא מיוחד.\n",
"עליו הוא לא מדבר ככה, אבל אני יודעת\n",
"יום יום. הוא עובר עם כולם ואני יודעת שהוא עדיין שם.\n",
"הוא לא קורא את המילים שאני כל כך שונאת.\n",
"אני תוהים אם הוא ילך עם העיניים שלי או עם המבט שלו.\n",
"הוא יושב שם, מחייך, לא מדבר\n",
"אבל אני לא שומעת אותו.\n",
"הוא צועק בקול רם מדי בשבילי.\n",
"הראש מסתובב אליו ושונא להיות לבד.\n",
"\n",
"-\n",
"היא הייתה יושבת לי ומחכה לי בשקיקה.\n",
"ובכל פעם שהייתי חוזרת חזרה לבית אני\n",
"כשאני מסיימת תיכון אני רואה את פניה והולכת לים\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"seed = 10000\n",
"model.seed = seed\n",
"generated_max_length = 1024\n",
"\n",
"#Changing temperature, top_k and top_p does not seem to change the outcome\n",
"outputs = model.generate(\n",
" input_ids = inputs[\"input_ids\"],\n",
" max_length=generated_max_length,\n",
" pad_token_id = model.config.eos_token_id,\n",
" prng_key=jax.random.PRNGKey(seed),\n",
" top_k=40,\n",
" top_p=0.92,\n",
" do_sample=True,\n",
" no_repeat_ngram_size=3)\n",
"\n",
"output_sequence = outputs['sequences'].squeeze(0)\n",
"text = tokenizer.decode(output_sequence, clean_up_tokenization_spaces=True)\n",
"\n",
"print(text)\n"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "luBCNTcZ5IPm",
"outputId": "98851ab5-aac6-459f-cb2b-b5140c56524c"
},
"execution_count": 4,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"\n",
"\n",
"\n",
"שם היצירה: ללא שם\n",
"מאת: נ. עומר\n",
"\n",
"ואולי רק תישארי איתי.\n",
"אז אני כאן, מחכה לתשובה,\n",
"ואני פה. מחכה לרכבת שתגיע,\n",
"ואת פה\n",
"אני מביט בעינייך,\n",
"שאת שם.\n",
"אני יודע,\n",
"אם לא תישאר, אם לא תלכי לי לדרכך,\n",
"אני פה, מחפש את החושך,\n",
"ולחיות חיים של אושר.\n",
"שזה יהיה,\n",
"אולי את שם, את שם\n",
"אני מחכה לראות אותך,\n",
"בכל יום ובכל שעה.\n",
"תוהה אם יש עוד יום או שלוש,\n",
"מחפש. מחפש.\n",
"עומד,\n",
"\n",
" \n",
"עיניו העצובות, העמוקות.\n",
"כעבור רגע.\n",
"הוא מתרחק,\n",
"הוא לא רוצה להיזכר בי,\n",
"אך כבר אינני מוצא.\n",
"אולי לא אראה אותו שוב,\n",
"אבל הוא רוצה להזכר בי, שוב.\n",
"למרות כל התחושות,\n",
"הם שואלים את השאלות\n",
"והוא שותק.\n",
"נמצא שם. הוא יושב,\n",
"את לא איתי. את לא איתי\n",
"אבל עכשיו, כבר לא יהיה.\n",
"כל פעם הוא נזכר בי, לא יודע, מה הוא חושב.\n",
"לא מרגיש.\n",
"\n",
"2\n",
"תפסתי את הרכבת הנכונה.\n",
"אמרתי: כן, אני זז.\n",
"ואחרי שנה, לא רציתי להיפרד.\n",
"רק את, היית איתי,\n",
"והרגשתי כל-כך טוב.\n",
"גם רציתי שיהיה עוד יום\n",
"לא האמנתי, שמחר יהיה יותר טוב. אבל לא האמניתי, זה\n",
"\n",
"3\n",
"\n",
"1\n",
"\n",
"9\n",
"6\n",
"5\n",
"6 היא הייתה כה קרובה, כל- כך יפה, כה מדהים.\n",
"6 אבל את לא יודעת כמה, כמה,\n",
"6 אולי הוא אמר את המשפט ההוא שלו,\n",
"7 הוא חייך אליך חיוך יפהפה,\n",
"1 היא הייתה יפה כמו השמש,\n",
"4 היא הייתה עדינה כמו השמש.\n",
"7 פעם היא הייתה כל כך יפה.\n",
"5 היום כבר לא ראיתי אותה מחייכת.\n",
"4 מחר זה יהיה היום האחרון.\n",
"8 פעם היא חזרה.\n",
"2 הוא אמר לה שהיום אני הולך לפגוש אותה.\n",
"1 פעם היא אמרה לו שהיום היא תראה איך הוא מעכל את זה.\n",
"9 פעם היא פגשה אותו,\n",
"8 היא הייתה מאוהבת,\n",
"3 אולי זה היה נכון, אבל עכשיו זה היה שונה.\n",
"3 פעם היא היתה כל כך מושלמת,\n",
"2 פעם היא שאלה אותו מה הוא מצא.\n",
"13 פעם היא סיפרה לו הכל.\n",
"12 פעם היא לא הייתה יפה, לא מאמינה.\n",
"21 פעם היא הלכה לפגוש אותו.\n",
"01 פעם היא ראתה אותו. כמה יפה. היא נראתה כל כך טוב. היא\n",
"6 פעם היא חייכה בפעם הראשונה,\n",
"01 אולי היא רצתה לומר משהו. אבל זה לא ממש הספיק לה.\n",
"71 פעם היא אהבה אותו. היא כל כך אהבה אותו, כל כך רצתה.\n",
"51 פעם היא רוותה אותו, יותר מידי, יותר מדי.\n",
"41 פעם היא חיממה לו את הלב,\n",
"71 אחד היא נתנה לו.\n",
"31 פעם היא נתנה את לבה, יותר מהכל.\n",
"53 פעם היא התאהבה בפעם הראשונה.\n",
"52 פעם היא יצאה לנשום אויר צח.\n",
"81 פעם היא ביקשה ממנו, יותר מכל.\n",
"18 פעם היא בכתה.\n",
"32 פעם היא קראה את הספר, יותר מ71 פעם.\n",
"10 פעם היא צחקה, יותר ויותר.\n",
"11 פעם היא למדה למבחנים.\n",
"25 פעם היא אפילו לא מצאה, אפילו לא שמה לב.\n",
"19 פעם היא חשבה, יותר נכון.\n",
"20 פעם היא החליטה, יותר מתמיד.\n",
"17 פעם היא חיה לבד.\n",
"26 פעם היא עברה, יותר לא.\n",
"22 פעם היא ניסתה להבין מה הוא עובר על מה שהוא כותב, יותר\n",
"19\n",
"\n",
"15 פעם היא התיישבה לידו.\n",
"15פעם היא בילתה איתו ביחד.\n",
"03 פעם היא ישבה ליד הבר, וחשבה.\n",
"14 פעם היא תפסה את הרכבת, פעם היא באה.\n",
"27 פעם היא הכירה אותו יותר.\n",
"24 פעם היא הגיעה לרכבת, הפעם זה היה ברכבת.\n",
"16 פעם היא הסתכלה בעיניו, יותר טוב מ-24 פעם. \n",
"181 פעם הוא הגיע לרכבת.\n",
"23 פעם היא עמדה לידו, כל פעם הוא התקרב אליה יותר. הוא נראה\n",
"17 שוב היא ישבת מולה.\n",
"30 פעם היא נחה לצידו.\n",
"07 פעם הוא ישב מולה, היא יושבת לידו ואומרת לו, יותר כיף\n",
"לה.  הידיים שלו עליה.\n",
"35 פעם הוא אמר, יותר חזק מזה הוא לא ראה.\n",
"04 פעם הוא ירד בתחנה, הוא הסתכל קדימה, יותר קרוב.\n",
"29 פעם היא המשיכה.\n",
"28 פעם היא עשתה את הדרך.\n",
"37 פעם היא עצרה, יותר לאט ממה שהוא ידע שהוא יהיה. \tהיא\n",
"היתה בטוחה שהיא מחכה, יותר זמן\n"
]
}
]
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment