Skip to content

Instantly share code, notes, and snippets.

@etrotta

etrotta/geminifileextractor.ipynb

Created May 8, 2024 10:50
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save etrotta/566da8c1e0e7a4110d7fede740644539 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save etrotta/566da8c1e0e7a4110d7fede740644539 to your computer and use it in GitHub Desktop.
Download ZIP
GeminiFileExtractor.ipynb
Display the source blob
Display the rendered blob
Raw
geminifileextractor.ipynb
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": [],
"authorship_tag": "ABX9TyOaA+Uo7gCeg1ORQ5tU0Rwh",
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
}
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/etrotta/566da8c1e0e7a4110d7fede740644539/geminifileextractor.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"source": [],
"metadata": {
"id": "P-JMcrBLuJvw"
}
},
{
"cell_type": "markdown",
"source": [
"## Overview - Major steps\n",
"- Clone (copy) this Notebook into your own Drive\n",
"- Go to <a href=\"https://aistudio.google.com/\">Google AI Studio</a> and log in with your Google account.</li>\n",
"- <a href=\"https://aistudio.google.com/app/apikey\">Create an API key</a>.</li>\n",
"- Save it as a Secret called \"GOOGLE_API_KEY\"\n",
"- Define the Data Models, Function Templates and Prompt (more on that later)\n",
"- Upload your files\n",
"\n",
"---\n",
"\n",
"## How it works\n",
"This Notebook allows for you to define Pydantic Models in order to format how the output will look like,\n",
"then upload your text files to Gemini so that it may extract data into a format that matches these Models.\n",
"\n",
"You should customize all three Configuration options - Data Models, Function Templates, and Prompt - as well as add your own Gemini API Key.\n",
"\n",
"The program will make one request to Gemini for each (File x Data Model) combination,\n",
" passing the entire file's contents each time, so be careful about your API Usage.\n",
"All functions are passed to Gemini, and it is up for Gemini to decide which ones to call.\n",
"\n",
"The model will only look at once file at a time -\n",
" it will not keep track of previously uploaded files nor information about previous prompts.\n",
" The prompt template must contain all information required, specially the file_text.\n",
"\n",
"---\n",
"\n",
"## Configuration:\n",
"### Template Functions\n",
"TEMPLATE_FUNCTIONS define the functions Gemini will try to call for each element it identifies in the submitted files.\n",
"Format: (function name, function description)\n",
"They will later be concatenated with each registered Model. For example:\n",
"```py\n",
"# Template functions\n",
"TEMPLATE_FUNCTIONS = [\n",
" (\"register_{}\", \"Register a new\"),\n",
"]\n",
"# Model\n",
"@register_model\n",
"class Character(pydantic.BaseModel):\n",
" \"important Character present in the book.\"\n",
" name: str = pydantic.Field(default=\"Unknown\", description=\"This character's real name.\")\n",
"```\n",
"Will generate the Gemini Tool Function\n",
"`register_character`: \"Register a new important Character present in the book.\"\n",
"\n",
"-----\n",
"Note: Despite we calling them Functions, they are not functions in the sense of python functions you have to create with `def`, they are just conceptually treated as functions Gemini tries to call, so that we can exploit the \"create the well-formatted arguments\" part to format data for us.\n",
"\n",
"-----\n",
"\n",
"### Prompt Template\n",
"The PROMPT_TEMPLATE will create the prompt actually passed to the model.\n",
"\n",
"The following keyword arguments will be passed to `PROMPT_TEMPLATE.format()`, which means you can include them in the template using {file_text} for example:\n",
"- `data_type`: The name of the Data Model it is being called for\n",
"- `file_text`: The contents of the file\n",
"---\n",
"\n",
"### Data Models\n",
"You have to create Pydantic models in order to define how the model output should be formatted.\n",
"They are converted to Gemini Functions via the `@register_model` decorator for simplicty.\n",
"\n",
"Currently the supported data types are\n",
"- str (String ; Any Text)\n",
"- bool (Boolean ; True or False)\n",
"- int (Integers ; Numbers with no fractional text)\n",
"- float (Decimal ; Numbers that may have a fractional part)\n",
"\n",
"And lists of the above types, but not nested lists, dictionaries nor nested models -\n",
"\n",
"- list[str], list[bool], list[int], list[float]\n",
"\n",
"And you can use `typing.Optional` to make them as nullable,\n",
"\n",
"- typing.Optional[str], typing.Optional[list[str]]\n",
"- typing.Optional[bool], typing.Optional[list[bool]]\n",
"- typing.Optional[int], typing.Optional[list[int]]\n",
"- typing.Optional[float], typing.Optional[list[float]]\n",
"\n",
"But you have to specify a `default` value or a `default_factory` function for it to be treated as non-required."
],
"metadata": {
"id": "ThSu79E0vHfg"
}
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"cellView": "form",
"id": "LPLx6Xubt6fW",
"outputId": "1cf3ea27-3ec5-4b45-d37b-d2d511ece98c"
},
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Requirement already satisfied: google-generativeai in /usr/local/lib/python3.10/dist-packages (0.5.2)\n",
"Requirement already satisfied: google-ai-generativelanguage==0.6.2 in /usr/local/lib/python3.10/dist-packages (from google-generativeai) (0.6.2)\n",
"Requirement already satisfied: google-api-core in /usr/local/lib/python3.10/dist-packages (from google-generativeai) (2.11.1)\n",
"Requirement already satisfied: google-api-python-client in /usr/local/lib/python3.10/dist-packages (from google-generativeai) (2.84.0)\n",
"Requirement already satisfied: google-auth>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from google-generativeai) (2.27.0)\n",
"Requirement already satisfied: protobuf in /usr/local/lib/python3.10/dist-packages (from google-generativeai) (3.20.3)\n",
"Requirement already satisfied: pydantic in /usr/local/lib/python3.10/dist-packages (from google-generativeai) (2.7.1)\n",
"Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from google-generativeai) (4.66.4)\n",
"Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from google-generativeai) (4.11.0)\n",
"Requirement already satisfied: proto-plus<2.0.0dev,>=1.22.3 in /usr/local/lib/python3.10/dist-packages (from google-ai-generativelanguage==0.6.2->google-generativeai) (1.23.0)\n",
"Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth>=2.15.0->google-generativeai) (5.3.3)\n",
"Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth>=2.15.0->google-generativeai) (0.4.0)\n",
"Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth>=2.15.0->google-generativeai) (4.9)\n",
"Requirement already satisfied: googleapis-common-protos<2.0.dev0,>=1.56.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core->google-generativeai) (1.63.0)\n",
"Requirement already satisfied: requests<3.0.0.dev0,>=2.18.0 in /usr/local/lib/python3.10/dist-packages (from google-api-core->google-generativeai) (2.31.0)\n",
"Requirement already satisfied: httplib2<1dev,>=0.15.0 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client->google-generativeai) (0.22.0)\n",
"Requirement already satisfied: google-auth-httplib2>=0.1.0 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client->google-generativeai) (0.1.1)\n",
"Requirement already satisfied: uritemplate<5,>=3.0.1 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client->google-generativeai) (4.1.1)\n",
"Requirement already satisfied: annotated-types>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from pydantic->google-generativeai) (0.6.0)\n",
"Requirement already satisfied: pydantic-core==2.18.2 in /usr/local/lib/python3.10/dist-packages (from pydantic->google-generativeai) (2.18.2)\n",
"Requirement already satisfied: grpcio<2.0dev,>=1.33.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core->google-generativeai) (1.63.0)\n",
"Requirement already satisfied: grpcio-status<2.0.dev0,>=1.33.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core->google-generativeai) (1.48.2)\n",
"Requirement already satisfied: pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 in /usr/local/lib/python3.10/dist-packages (from httplib2<1dev,>=0.15.0->google-api-python-client->google-generativeai) (3.1.2)\n",
"Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth>=2.15.0->google-generativeai) (0.6.0)\n",
"Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0.dev0,>=2.18.0->google-api-core->google-generativeai) (3.3.2)\n",
"Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0.dev0,>=2.18.0->google-api-core->google-generativeai) (3.7)\n",
"Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0.dev0,>=2.18.0->google-api-core->google-generativeai) (2.0.7)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0.dev0,>=2.18.0->google-api-core->google-generativeai) (2024.2.2)\n"
]
}
],
"source": [
"#@title Install dependencies and import libraries\n",
"!pip install google-generativeai\n",
"\n",
"import re\n",
"import typing\n",
"import pydantic\n",
"import time\n",
"from pydantic.fields import FieldInfo\n",
"from google.colab import files, userdata\n",
"import google.ai.generativelanguage as glm\n",
"import google.generativeai as genai"
]
},
{
"cell_type": "code",
"source": [
"#@title Configuration Variables [Requires Changes]\n",
"\"\"\"\n",
"You have to edit the PROMPT_TEMPLATE and TEMPLATE_FUNCTIONS to adjust to your use case.\n",
"For the rest, the default values should be OK\n",
"\n",
"If anything, you can try changing the Gemini version being used or change the GenerationConfig.\n",
"\"\"\"\n",
"\n",
"registered_models: dict[str, typing.Type[pydantic.BaseModel]] = {} # Pydantic Data Model Name => Model\n",
"declared_functions: dict[str, list[glm.FunctionDeclaration]] = {} # Pydantic Data Model Name => Function\n",
"\n",
"genai.configure(api_key=userdata.get('GOOGLE_API_KEY'))\n",
"\n",
"gemini = genai.GenerativeModel(\n",
" \"gemini-1.0-pro\",\n",
" generation_config=genai.GenerationConfig(temperature=0.4),\n",
")\n",
"\n",
"TEMPLATE_FUNCTIONS = [\n",
" (\"register_{}\", \"Register a new\"),\n",
" # (\"start_{}_event\", \"Add information about a new event related to a\"),\n",
"]\n",
"\n",
"PROMPT_TEMPLATE = \"\"\"Given the following book page, \\\n",
"identify all notable {data_type}s \\\n",
"and register each of them through Tools, \\\n",
"including all relevant information about them.\n",
"```\n",
"{file_text}\n",
"```\"\"\"\n",
"\n",
"TOOL_CONFIG = {\"function_calling_config\": \"ANY\"} # I recommend not touching this\n",
"FILE_ENCODING = \"UTF-8\"\n"
],
"metadata": {
"id": "iQtymD9Pxzp1"
},
"execution_count": 2,
"outputs": []
},
{
"cell_type": "code",
"source": [
"#@title Define the code to convert Data Models into Gemini Functions\n",
"\"\"\"Most technical Cell of the entire Notebook - Formats the pydantic models as Gemini Functions.\n",
"Do not worry about understanding it, just run and move on to the next Cell\n",
"\"\"\"\n",
"\n",
"class SchemaConfig(typing.TypedDict):\n",
" description: str\n",
" nullable: bool\n",
"\n",
"\n",
"# def _format_type(annotation: type, **kwargs: typing.Unpack[SchemaConfig]) -> glm.Schema: # Typing for (3.11+? 3.12+?)\n",
"def _format_type(annotation: type, **kwargs) -> glm.Schema:\n",
" mapping = {\n",
" str: glm.Schema(type=glm.Type.STRING, **kwargs),\n",
" bool: glm.Schema(type=glm.Type.BOOLEAN, **kwargs),\n",
" int: glm.Schema(type=glm.Type.INTEGER, **kwargs),\n",
" float: glm.Schema(type=glm.Type.NUMBER, **kwargs),\n",
" list[str]: glm.Schema(type=glm.Type.ARRAY, items=glm.Schema(type=glm.Type.STRING), **kwargs),\n",
" list[bool]: glm.Schema(type=glm.Type.ARRAY, items=glm.Schema(type=glm.Type.BOOLEAN), **kwargs),\n",
" list[int]: glm.Schema(type=glm.Type.ARRAY, items=glm.Schema(type=glm.Type.INTEGER), **kwargs),\n",
" list[float]: glm.Schema(type=glm.Type.ARRAY, items=glm.Schema(type=glm.Type.NUMBER), **kwargs),\n",
" }\n",
" return mapping[annotation]\n",
"\n",
"\n",
"def format_parameter(field: FieldInfo) -> glm.Schema:\n",
" assert field.annotation is not None, f\"Field {field} missing annotation\"\n",
" assert field.description is not None, f\"Field {field} missing description\"\n",
" _optional_mappings = {typing.Optional[ann]: ann for type_ in (str, bool, int, float) for ann in (type_, list[type_])}\n",
" annotation = field.annotation\n",
" nullable = annotation in _optional_mappings\n",
" if nullable:\n",
" annotation = _optional_mappings[annotation]\n",
" config = SchemaConfig(description=field.description, nullable=nullable)\n",
" return _format_type(annotation, **config)\n",
"\n",
"\n",
"# def register_model[ModelType: typing.Type[pydantic.BaseModel]](model: ModelType) -> ModelType: # 3.12+\n",
"def register_model(model):\n",
" \"\"\"Register the Model and declare Tool functions to manage instances of it.\"\"\"\n",
" snake_case_name = '_'.join(re.findall(r\"([A-Z][a-z0-9\\_]+)\", model.__name__)).lower()\n",
"\n",
" registered_models[model.__name__] = model\n",
"\n",
" parameters: dict[str, glm.Schema] = {}\n",
" required: list[str] = []\n",
" for key, field in model.model_fields.items():\n",
" parameters[key] = format_parameter(field)\n",
" if field.is_required():\n",
" required.append(key)\n",
"\n",
" model_schema = glm.Schema(\n",
" type=glm.Type.OBJECT,\n",
" nullable=False,\n",
" properties=parameters,\n",
" required=required,\n",
" )\n",
" functions = {\n",
" model.__name__: [glm.FunctionDeclaration(\n",
" name=function_name.format(snake_case_name),\n",
" description=f\"{function_description} {model.__doc__}\",\n",
" parameters=model_schema,\n",
" )\n",
" for function_name, function_description in TEMPLATE_FUNCTIONS]\n",
" }\n",
" declared_functions.update(functions)\n",
" return model\n",
"\n",
"def clear_models_and_functions() -> None:\n",
" \"Unregister all models and functions previously declared by `register_model`\"\n",
" registered_models.clear()\n",
" declared_functions.clear()\n",
"\n",
"def collect_models() -> dict[str, typing.Type[pydantic.BaseModel]]:\n",
" \"Lists all models registered as `@register_model`s\"\n",
" return registered_models.copy()\n",
"\n",
"def collect_tools(model_name: typing.Optional[str] = None) -> glm.Tool:\n",
" \"Collects functions declared by `@register_model`, optionally filtering by the relevant model\"\n",
" if model_name:\n",
" functions = declared_functions[model_name]\n",
" else:\n",
" functions = [func for sublist in declared_functions.values() for func in sublist]\n",
" return glm.Tool(function_declarations=functions)\n",
"\n",
"def get_model_from_function_name(function_name: str) -> typing.Type[pydantic.BaseModel]:\n",
" owners = {function.name: owner for owner, functions in declared_functions.items() for function in functions}\n",
" return registered_models[owners[function_name]]\n"
],
"metadata": {
"cellView": "form",
"id": "1C2IQpCmwb3t"
},
"execution_count": 3,
"outputs": []
},
{
"cell_type": "code",
"source": [
"#@title Creater your Data Models [Requires Changes]\n",
"\"\"\"You must edit this Cell.\n",
"It defines which models you will extract out of the files you upload in the end.\n",
"\"\"\"\n",
"\n",
"@register_model\n",
"class Character(pydantic.BaseModel):\n",
" \"\"\"important Character present in the book.\"\"\"\n",
"\n",
" name: str = pydantic.Field(default=\"Unknown\", description=\"This character's real name.\")\n",
" description: str = pydantic.Field(\n",
" description=\"A short description including the most important details about this character.\"\n",
" )\n",
" abilities: list[str] = pydantic.Field(description=\"What this character is good at.\")\n"
],
"metadata": {
"id": "lR3hkjsOwcn0"
},
"execution_count": 4,
"outputs": []
},
{
"cell_type": "code",
"source": [
"#@title Create the function that actually calls the Gemini API and formats its output\n",
"\"\"\"Define the actual python function that calls Gemini for a given (Model X File) combination\"\"\"\n",
"\n",
"def run_for_file(model, file_name: str, file_content: str, data_model: type[pydantic.BaseModel]):\n",
" assert \"file_text\" in PROMPT_TEMPLATE, \"\"\"ERROR: You seem to have forgotten to include {file_text} in the Prompt Template.\n",
" Gemini would not have access to the file's contents this way. If this was intentional, feel free to delete this `assert` line.\"\"\"\n",
"\n",
" prompt = PROMPT_TEMPLATE.format(file_text=file_content, data_type=data_model.__name__)\n",
" model_tools = collect_tools(data_model.__name__)\n",
" response = model.generate_content(\n",
" prompt,\n",
" tools=model_tools,\n",
" tool_config=TOOL_CONFIG,\n",
" )\n",
"\n",
" # Just logging / for assisting in in debugging\n",
" for i, candidate in enumerate(response.candidates, 1):\n",
" for j, part in enumerate(candidate.content.parts, 1):\n",
" print(f\"File {file_name!r} - {data_model.__name__} Data\"\n",
" f\" Candidate {i}/{len(response.candidates)} Part {j}/{len(candidate.content.parts)}:\")\n",
" if part.text:\n",
" # *Should* always be empty. In practice it can not be thanks to bugs on their side.\n",
" print(f\"- Output Text: {part.text!r}\")\n",
" if part.function_call.name:\n",
" _function_args = dict(part.function_call.args.items()) if part.function_call.args else None\n",
" print(f\"- Output Function Call: {part.function_call.name!r}({_function_args!r})\")\n",
"\n",
"\n",
" result = {}\n",
"\n",
" for part in (part for candidate in response.candidates for part in candidate.content.parts):\n",
" call = part.function_call\n",
" if call.name == \"\" or call.args is None:\n",
" continue\n",
" arguments = dict(call.args.items())\n",
"\n",
" try:\n",
" model_type = get_model_from_function_name(call.name)\n",
" except KeyError:\n",
" print(f\"The model tried to call a function that does not exists: {call.name}. Skipping this Part.\")\n",
" continue\n",
" if model_type is not data_model:\n",
" print(f\"Expected for model to operate on a {data_model}, instead model is operating on a {model_type}. Skipping this Part\")\n",
" # ^ Extremely unlikely - we only give information about one model at a time,\n",
" # it would have to randomly pick a string correspondent to different existing model by coincidence...\n",
" continue\n",
"\n",
" try:\n",
" created_model = model_type.model_validate(arguments) # Feel free to change this logic a bit, e.g. specifying `strict=...`\n",
" except pydantic.ValidationError:\n",
" print(f\"Validation Error trying to load {model_type} with arguments {arguments!r}\")\n",
" continue\n",
" result.setdefault(model_type.__name__, []).append(created_model)\n",
" return result\n"
],
"metadata": {
"cellView": "form",
"id": "sBwOcPUOweVW"
},
"execution_count": 5,
"outputs": []
},
{
"cell_type": "code",
"source": [
"#@title Upload your files and runs the above python function for each of them\n",
"uploaded = files.upload()\n",
"\n",
"complete_output = {}\n",
"for file_name, file_content in uploaded.items():\n",
" file_content = file_content.decode(FILE_ENCODING)\n",
" print(f\"File preview: {file_name}\\n--------------\\n{file_content[:100]}\\n--------------\")\n",
"\n",
" file_results: dict[str, list] = {}\n",
" for data_model in collect_models().values():\n",
" print(f\"Processing file {file_name!r} for {data_model.__name__}\")\n",
" result = run_for_file(gemini, file_name=file_name, file_content=file_content, data_model=data_model)\n",
" print(result)\n",
" file_results.update(result)\n",
" print(\"Sleeping for 4 seconds in consideration of rate limits\")\n",
" time.sleep(4)\n",
"\n",
" complete_output[file_name] = file_results\n"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "ii5ihhHjwfkm",
"outputId": "64356e3e-db5d-44e0-a6e1-63585c1f32fe"
},
"execution_count": 6,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"<IPython.core.display.HTML object>"
],
"text/html": [
"\n",
" <input type=\"file\" id=\"files-d96781a0-dd17-44b2-9a56-25f08e164346\" name=\"files[]\" multiple disabled\n",
" style=\"border:none\" />\n",
" <output id=\"result-d96781a0-dd17-44b2-9a56-25f08e164346\">\n",
" Upload widget is only available when the cell has been executed in the\n",
" current browser session. Please rerun this cell to enable.\n",
" </output>\n",
" <script>// Copyright 2017 Google LLC\n",
"//\n",
"// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"// you may not use this file except in compliance with the License.\n",
"// You may obtain a copy of the License at\n",
"//\n",
"// http://www.apache.org/licenses/LICENSE-2.0\n",
"//\n",
"// Unless required by applicable law or agreed to in writing, software\n",
"// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"// See the License for the specific language governing permissions and\n",
"// limitations under the License.\n",
"\n",
"/**\n",
" * @fileoverview Helpers for google.colab Python module.\n",
" */\n",
"(function(scope) {\n",
"function span(text, styleAttributes = {}) {\n",
" const element = document.createElement('span');\n",
" element.textContent = text;\n",
" for (const key of Object.keys(styleAttributes)) {\n",
" element.style[key] = styleAttributes[key];\n",
" }\n",
" return element;\n",
"}\n",
"\n",
"// Max number of bytes which will be uploaded at a time.\n",
"const MAX_PAYLOAD_SIZE = 100 * 1024;\n",
"\n",
"function _uploadFiles(inputId, outputId) {\n",
" const steps = uploadFilesStep(inputId, outputId);\n",
" const outputElement = document.getElementById(outputId);\n",
" // Cache steps on the outputElement to make it available for the next call\n",
" // to uploadFilesContinue from Python.\n",
" outputElement.steps = steps;\n",
"\n",
" return _uploadFilesContinue(outputId);\n",
"}\n",
"\n",
"// This is roughly an async generator (not supported in the browser yet),\n",
"// where there are multiple asynchronous steps and the Python side is going\n",
"// to poll for completion of each step.\n",
"// This uses a Promise to block the python side on completion of each step,\n",
"// then passes the result of the previous step as the input to the next step.\n",
"function _uploadFilesContinue(outputId) {\n",
" const outputElement = document.getElementById(outputId);\n",
" const steps = outputElement.steps;\n",
"\n",
" const next = steps.next(outputElement.lastPromiseValue);\n",
" return Promise.resolve(next.value.promise).then((value) => {\n",
" // Cache the last promise value to make it available to the next\n",
" // step of the generator.\n",
" outputElement.lastPromiseValue = value;\n",
" return next.value.response;\n",
" });\n",
"}\n",
"\n",
"/**\n",
" * Generator function which is called between each async step of the upload\n",
" * process.\n",
" * @param {string} inputId Element ID of the input file picker element.\n",
" * @param {string} outputId Element ID of the output display.\n",
" * @return {!Iterable<!Object>} Iterable of next steps.\n",
" */\n",
"function* uploadFilesStep(inputId, outputId) {\n",
" const inputElement = document.getElementById(inputId);\n",
" inputElement.disabled = false;\n",
"\n",
" const outputElement = document.getElementById(outputId);\n",
" outputElement.innerHTML = '';\n",
"\n",
" const pickedPromise = new Promise((resolve) => {\n",
" inputElement.addEventListener('change', (e) => {\n",
" resolve(e.target.files);\n",
" });\n",
" });\n",
"\n",
" const cancel = document.createElement('button');\n",
" inputElement.parentElement.appendChild(cancel);\n",
" cancel.textContent = 'Cancel upload';\n",
" const cancelPromise = new Promise((resolve) => {\n",
" cancel.onclick = () => {\n",
" resolve(null);\n",
" };\n",
" });\n",
"\n",
" // Wait for the user to pick the files.\n",
" const files = yield {\n",
" promise: Promise.race([pickedPromise, cancelPromise]),\n",
" response: {\n",
" action: 'starting',\n",
" }\n",
" };\n",
"\n",
" cancel.remove();\n",
"\n",
" // Disable the input element since further picks are not allowed.\n",
" inputElement.disabled = true;\n",
"\n",
" if (!files) {\n",
" return {\n",
" response: {\n",
" action: 'complete',\n",
" }\n",
" };\n",
" }\n",
"\n",
" for (const file of files) {\n",
" const li = document.createElement('li');\n",
" li.append(span(file.name, {fontWeight: 'bold'}));\n",
" li.append(span(\n",
" `(${file.type || 'n/a'}) - ${file.size} bytes, ` +\n",
" `last modified: ${\n",
" file.lastModifiedDate ? file.lastModifiedDate.toLocaleDateString() :\n",
" 'n/a'} - `));\n",
" const percent = span('0% done');\n",
" li.appendChild(percent);\n",
"\n",
" outputElement.appendChild(li);\n",
"\n",
" const fileDataPromise = new Promise((resolve) => {\n",
" const reader = new FileReader();\n",
" reader.onload = (e) => {\n",
" resolve(e.target.result);\n",
" };\n",
" reader.readAsArrayBuffer(file);\n",
" });\n",
" // Wait for the data to be ready.\n",
" let fileData = yield {\n",
" promise: fileDataPromise,\n",
" response: {\n",
" action: 'continue',\n",
" }\n",
" };\n",
"\n",
" // Use a chunked sending to avoid message size limits. See b/62115660.\n",
" let position = 0;\n",
" do {\n",
" const length = Math.min(fileData.byteLength - position, MAX_PAYLOAD_SIZE);\n",
" const chunk = new Uint8Array(fileData, position, length);\n",
" position += length;\n",
"\n",
" const base64 = btoa(String.fromCharCode.apply(null, chunk));\n",
" yield {\n",
" response: {\n",
" action: 'append',\n",
" file: file.name,\n",
" data: base64,\n",
" },\n",
" };\n",
"\n",
" let percentDone = fileData.byteLength === 0 ?\n",
" 100 :\n",
" Math.round((position / fileData.byteLength) * 100);\n",
" percent.textContent = `${percentDone}% done`;\n",
"\n",
" } while (position < fileData.byteLength);\n",
" }\n",
"\n",
" // All done.\n",
" yield {\n",
" response: {\n",
" action: 'complete',\n",
" }\n",
" };\n",
"}\n",
"\n",
"scope.google = scope.google || {};\n",
"scope.google.colab = scope.google.colab || {};\n",
"scope.google.colab._files = {\n",
" _uploadFiles,\n",
" _uploadFilesContinue,\n",
"};\n",
"})(self);\n",
"</script> "
]
},
"metadata": {}
},
{
"output_type": "stream",
"name": "stdout",
"text": [
"Saving N3702IR-10.txt to N3702IR-10.txt\n",
"Saving N3702IR-11.txt to N3702IR-11.txt\n",
"Saving N3702IR-12.txt to N3702IR-12.txt\n",
"Saving N3702IR-13.txt to N3702IR-13.txt\n",
"Saving N3702IR-14.txt to N3702IR-14.txt\n",
"File preview: N3702IR-10.txt\n",
"--------------\n",
"10話 シロクロ\r\n",
"\r\n",
"\r\n",
"　スープをお替りした2人は、それもあっという間に飲み干していた。匂いを漏らさぬように張っていた風魔法の結界を解除する。\r\n",
"\r\n",
"　これ、料理中に出る煙なんかも遮断できるから\n",
"--------------\n",
"Processing file 'N3702IR-10.txt' for Character\n",
"File 'N3702IR-10.txt' - Character Data Candidate 1/1 Part 1/2:\n",
"- Output Function Call: 'register_character'({'name': 'シロ', 'description': '金髪の猫獣人の少女。ワクワクした表情でトールを見つめる。', 'abilities': ['猫の力']})\n",
"File 'N3702IR-10.txt' - Character Data Candidate 1/1 Part 2/2:\n",
"- Output Function Call: 'register_character'({'description': '銀髪の犬獣人の少女。眠たげな表情でトールを見つめる。', 'abilities': ['犬の力'], 'name': 'クロ'})\n",
"{'Character': [Character(name='シロ', description='金髪の猫獣人の少女。ワクワクした表情でトールを見つめる。', abilities=['猫の力']), Character(name='クロ', description='銀髪の犬獣人の少女。眠たげな表情でトールを見つめる。', abilities=['犬の力'])]}\n",
"Sleeping for 4 seconds in consideration of rate limits\n",
"File preview: N3702IR-11.txt\n",
"--------------\n",
"11話 魔法練習\n",
"\n",
"「うにゅぅ」\n",
"「わふぅ」\n",
"\n",
"　シロとクロが、マットの上で微睡んでいる。それを見ながら、俺は今後のことを考えていた。\n",
"\n",
"　2人を今後とも保護し、育てていくのは決定事項\n",
"--------------\n",
"Processing file 'N3702IR-11.txt' for Character\n",
"File 'N3702IR-11.txt' - Character Data Candidate 1/1 Part 1/3:\n",
"- Output Function Call: 'register_character'({'abilities': ['Magic', 'Cooking'], 'description': 'A kind and determined man who has taken in two young girls, Shiro and Kuro, and is teaching them magic.', 'name': 'Tohl'})\n",
"File 'N3702IR-11.txt' - Character Data Candidate 1/1 Part 2/3:\n",
"- Output Function Call: 'register_character'({'name': 'Shiro', 'description': 'A young girl who is eager to learn magic and become strong.', 'abilities': ['Magic']})\n",
"File 'N3702IR-11.txt' - Character Data Candidate 1/1 Part 3/3:\n",
"- Output Function Call: 'register_character'({'description': 'A young girl who is also eager to learn magic and become strong.', 'abilities': ['Magic'], 'name': 'Kuro'})\n",
"{'Character': [Character(name='Tohl', description='A kind and determined man who has taken in two young girls, Shiro and Kuro, and is teaching them magic.', abilities=['Magic', 'Cooking']), Character(name='Shiro', description='A young girl who is eager to learn magic and become strong.', abilities=['Magic']), Character(name='Kuro', description='A young girl who is also eager to learn magic and become strong.', abilities=['Magic'])]}\n",
"Sleeping for 4 seconds in consideration of rate limits\n",
"File preview: N3702IR-12.txt\n",
"--------------\n",
"12話 下水暮らし\n",
"\n",
"\n",
"　クロシロを助けてから7日。\n",
"\n",
"　俺は2人を連れて、下水の奥に入り込んでいた。\n",
"\n",
"「クロ、どうだ？」\n",
"「あっちから蟲のにおいー」\n",
"「よし、シロはそっちから追\n",
"--------------\n",
"Processing file 'N3702IR-12.txt' for Character\n",
"File 'N3702IR-12.txt' - Character Data Candidate 1/1 Part 1/2:\n",
"- Output Function Call: 'register_character'({'name': 'クロ', 'description': '猫獣人。敏捷性と腕力に優れる。爪と風魔法が得意。', 'abilities': ['爪', '風魔法']})\n",
"File 'N3702IR-12.txt' - Character Data Candidate 1/1 Part 2/2:\n",
"- Output Function Call: 'register_character'({'abilities': ['闇魔法', '火魔法'], 'name': 'シロ', 'description': '猫獣人。魔法への才能がある。闇魔法と火魔法が得意。'})\n",
"{'Character': [Character(name='クロ', description='猫獣人。敏捷性と腕力に優れる。爪と風魔法が得意。', abilities=['爪', '風魔法']), Character(name='シロ', description='猫獣人。魔法への才能がある。闇魔法と火魔法が得意。', abilities=['闇魔法', '火魔法'])]}\n",
"Sleeping for 4 seconds in consideration of rate limits\n",
"File preview: N3702IR-13.txt\n",
"--------------\n",
"13話 下水の外\n",
"\n",
"\n",
"　ザワザワ。\n",
"\n",
"　俺は町の雑踏の中にいた。場所は、複数の露店がが立ち並ぶ小さな広場だ。\n",
"\n",
"　目立ってないよな？　心なしか、皆の視線が俺に向いている気がするが、自意\n",
"--------------\n",
"Processing file 'N3702IR-13.txt' for Character\n",
"File 'N3702IR-13.txt' - Character Data Candidate 1/1 Part 1/1:\n",
"- Output Function Call: 'register_character'({'abilities': [], 'name': 'Unnamed Man', 'description': 'A bald man with a potbelly who attacks the protagonist for allegedly stealing bread.'})\n",
"{'Character': [Character(name='Unnamed Man', description='A bald man with a potbelly who attacks the protagonist for allegedly stealing bread.', abilities=[])]}\n",
"Sleeping for 4 seconds in consideration of rate limits\n",
"File preview: N3702IR-14.txt\n",
"--------------\n",
"14話 カロリナ\n",
"\n",
"\n",
"　声をかけた女性が、俺を見つめる。\n",
"\n",
"　バッサリと短くしてある髪は、汚れているが元々は赤いんだろう。肌も埃塗れだが、多分元々は黄色人種系かな？　年齢は20は超えていそ\n",
"--------------\n",
"Processing file 'N3702IR-14.txt' for Character\n",
"File 'N3702IR-14.txt' - Character Data Candidate 1/1 Part 1/1:\n",
"- Output Function Call: 'register_character'({'abilities': ['本の写本作成'], 'description': '天竜の火災で全身に火傷を負い、視力と利き手を失った女性。現在はスラムで生活している。', 'name': 'カロリナ'})\n",
"{'Character': [Character(name='カロリナ', description='天竜の火災で全身に火傷を負い、視力と利き手を失った女性。現在はスラムで生活している。', abilities=['本の写本作成'])]}\n",
"Sleeping for 4 seconds in consideration of rate limits\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"#@title Writes the output to a file and downloads it\n",
"adapter = pydantic.TypeAdapter(dict[str, dict[str, list]])\n",
"\n",
"with open(\"results.json\", \"wb\") as out_file:\n",
" out_file.write(adapter.dump_json(complete_output, indent=4))\n",
"\n",
"files.download(\"results.json\")\n"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "ybmkmPqawgFF",
"outputId": "106fc634-885f-40c8-ec2a-238387270b5c"
},
"execution_count": 8,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"<IPython.core.display.Javascript object>"
],
"application/javascript": [
"\n",
" async function download(id, filename, size) {\n",
" if (!google.colab.kernel.accessAllowed) {\n",
" return;\n",
" }\n",
" const div = document.createElement('div');\n",
" const label = document.createElement('label');\n",
" label.textContent = `Downloading \"${filename}\": `;\n",
" div.appendChild(label);\n",
" const progress = document.createElement('progress');\n",
" progress.max = size;\n",
" div.appendChild(progress);\n",
" document.body.appendChild(div);\n",
"\n",
" const buffers = [];\n",
" let downloaded = 0;\n",
"\n",
" const channel = await google.colab.kernel.comms.open(id);\n",
" // Send a message to notify the kernel that we're ready.\n",
" channel.send({})\n",
"\n",
" for await (const message of channel.messages) {\n",
" // Send a message to notify the kernel that we're ready.\n",
" channel.send({})\n",
" if (message.buffers) {\n",
" for (const buffer of message.buffers) {\n",
" buffers.push(buffer);\n",
" downloaded += buffer.byteLength;\n",
" progress.value = downloaded;\n",
" }\n",
" }\n",
" }\n",
" const blob = new Blob(buffers, {type: 'application/binary'});\n",
" const a = document.createElement('a');\n",
" a.href = window.URL.createObjectURL(blob);\n",
" a.download = filename;\n",
" div.appendChild(a);\n",
" a.click();\n",
" div.remove();\n",
" }\n",
" "
]
},
"metadata": {}
},
{
"output_type": "display_data",
"data": {
"text/plain": [
"<IPython.core.display.Javascript object>"
],
"application/javascript": [
"download(\"download_5cfc97f2-5d08-41ba-bbd2-68eaff3d6a47\", \"results.json\", 2718)"
]
},
"metadata": {}
}
]
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment