ProTransGenerative.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "pAY6WxodCFiY"
      },
      "outputs": [],
      "source": [
        "!pip install -q transformers"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from transformers import XLNetLMHeadModel, XLNetTokenizer,pipeline\n",
        "import re\n",
        "import os\n",
        "import requests\n",
        "from tqdm.auto import tqdm"
      ],
      "metadata": {
        "id": "r4lAnlvwCL4P"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = XLNetTokenizer.from_pretrained(\"Rostlab/prot_xlnet\", do_lower_case=False)"
      ],
      "metadata": {
        "id": "Qq1rNceGCMm4"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = XLNetLMHeadModel.from_pretrained(\"Rostlab/prot_xlnet\")"
      ],
      "metadata": {
        "id": "frhWA9KPCOYj"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Ensure sequences_Example is properly tokenized\n",
        "sequences_Example = \"A E T C Z A O\"\n",
        "sequences_Example = re.sub(r\"[UZOB]\", \"<unk>\", sequences_Example)"
      ],
      "metadata": {
        "id": "pJDWYvxpCQUS"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ids = tokenizer.encode(sequences_Example, add_special_tokens=False)"
      ],
      "metadata": {
        "id": "ovL0urFSCZRI"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "filtered_ids = [token_id for token_id in ids if token_id < tokenizer.vocab_size]\n",
        "if not filtered_ids:\n",
        "    filtered_ids = [tokenizer.unk_token_id]\n",
        "input_ids = torch.tensor(filtered_ids).unsqueeze(0)"
      ],
      "metadata": {
        "id": "NLRVIARKOTxt"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "max_length = 100\n",
        "temperature = 1.0\n",
        "k = 0\n",
        "p = 0.9\n",
        "repetition_penalty = 1.0\n",
        "num_return_sequences = 3"
      ],
      "metadata": {
        "id": "rwMXK06CD3yz"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "output_ids = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    max_length=max_length,\n",
        "    temperature=temperature,\n",
        "    top_k=k,\n",
        "    top_p=p,\n",
        "    repetition_penalty=repetition_penalty,\n",
        "    do_sample=True,\n",
        "    num_return_sequences=num_return_sequences,\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MeIRCQkoD6XT",
        "outputId": "390a8b61-dca5-48ba-81b7-fd2a57ac1b35"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "This is a friendly reminder - the current text generation call will exceed the model's predefined maximum length (-1). Depending on the model, you may observe exceptions, performance degradation, or nothing at all.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Decode the generated sequences\n",
        "output_sequences = [tokenizer.decode(output_id, skip_special_tokens=True) for output_id in output_ids]\n"
      ],
      "metadata": {
        "id": "v6dPhIRDEBW7"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Print generated sequences\n",
        "print('Generated Sequences\\n')\n",
        "for output_sequence in output_sequences:\n",
        "    print(output_sequence)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LrDYNA0jEI0u",
        "outputId": "112b5586-cf68-42ed-b27c-fe6481f0c214"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated Sequences\n",
            "\n",
            "A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C\n",
            "A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C A A A P A E T C\n",
            "A E T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C A A A P E Q T C\n"
          ]
        }
      ]
    }
  ]
}