novelai-python 0.5.0__tar.gz → 0.5.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (80) hide show
  1. {novelai_python-0.5.0 → novelai_python-0.5.2}/PKG-INFO +23 -4
  2. {novelai_python-0.5.0 → novelai_python-0.5.2}/README.md +21 -3
  3. {novelai_python-0.5.0 → novelai_python-0.5.2}/pyproject.toml +2 -1
  4. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/_cost.py +2 -2
  5. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate/__init__.py +27 -8
  6. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate/_enum.py +27 -24
  7. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate/_schema.py +20 -25
  8. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate_stream.py +1 -1
  9. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tokenizer/__init__.py +7 -0
  10. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tokenizer/clip_simple_tokenizer.py +8 -0
  11. novelai_python-0.5.2/src/novelai_python/tool/image_metadata/__init__.py +297 -0
  12. novelai_python-0.5.2/src/novelai_python/tool/image_metadata/_type.py +21 -0
  13. novelai_python-0.5.2/src/novelai_python/tool/image_metadata/bch_utils.py +160 -0
  14. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tool/image_metadata/lsb_extractor.py +22 -19
  15. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/utils/__init__.py +1 -1
  16. novelai_python-0.5.0/src/novelai_python/tool/image_metadata/__init__.py +0 -225
  17. {novelai_python-0.5.0 → novelai_python-0.5.2}/LICENSE +0 -0
  18. {novelai_python-0.5.0 → novelai_python-0.5.2}/NOTICE.MD +0 -0
  19. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/__init__.py +0 -0
  20. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_enum.py +0 -0
  21. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_exceptions.py +0 -0
  22. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/__init__.py +0 -0
  23. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/ai/__init__.py +0 -0
  24. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/ai/generate.py +0 -0
  25. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/ai/generate_image.py +0 -0
  26. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/ai/generate_stream.py +0 -0
  27. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/ai/generate_voice.py +0 -0
  28. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/ai/upscale.py +0 -0
  29. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/schema.py +0 -0
  30. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/user/__init__.py +0 -0
  31. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/user/information.py +0 -0
  32. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/user/login.py +0 -0
  33. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/_response/user/subscription.py +0 -0
  34. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/credential/ApiToken.py +0 -0
  35. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/credential/JwtToken.py +0 -0
  36. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/credential/UserAuth.py +0 -0
  37. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/credential/__init__.py +0 -0
  38. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/credential/_base.py +0 -0
  39. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/__init__.py +0 -0
  40. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/LICENSE +0 -0
  41. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/__init__.py +0 -0
  42. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/_const.py +0 -0
  43. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/_enum.py +0 -0
  44. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/augment_image/__init__.py +0 -0
  45. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/augment_image/_enum.py +0 -0
  46. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate/_const.py +0 -0
  47. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate/themes.json +0 -0
  48. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate_image/__init__.py +0 -0
  49. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate_image/suggest_tags.py +0 -0
  50. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate_voice/__init__.py +0 -0
  51. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/generate_voice/_enum.py +0 -0
  52. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/ai/upscale.py +0 -0
  53. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/schema.py +0 -0
  54. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/user/__init__.py +0 -0
  55. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/user/information.py +0 -0
  56. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/user/login.py +0 -0
  57. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/sdk/user/subscription.py +0 -0
  58. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/server.py +0 -0
  59. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tokenizer/nerdstash_v1.model +0 -0
  60. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tokenizer/nerdstash_v2.model +0 -0
  61. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tokenizer/novelai.model +0 -0
  62. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tokenizer/novelai_v2.model +0 -0
  63. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tool/__init__.py +0 -0
  64. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tool/image_metadata/lsb_injector.py +0 -0
  65. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tool/paint_mask/__init__.py +0 -0
  66. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tool/random_prompt/__init__.py +0 -0
  67. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tool/random_prompt/tag.py +0 -0
  68. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tool/random_prompt/tag_character.py +0 -0
  69. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/tool/random_prompt/tag_nsfw.py +0 -0
  70. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/utils/encode.py +0 -0
  71. {novelai_python-0.5.0 → novelai_python-0.5.2}/src/novelai_python/utils/useful.py +0 -0
  72. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/__init__.py +0 -0
  73. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/test_generate_voice.py +0 -0
  74. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/test_random_prompt.py +0 -0
  75. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/test_server.py +0 -0
  76. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/test_server_run.py +0 -0
  77. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/test_upscale.py +0 -0
  78. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/test_user_information.py +0 -0
  79. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/test_user_login.py +0 -0
  80. {novelai_python-0.5.0 → novelai_python-0.5.2}/tests/test_user_subscription.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: novelai-python
3
- Version: 0.5.0
3
+ Version: 0.5.2
4
4
  Summary: NovelAI Python Binding With Pydantic
5
5
  Author-Email: sudoskys <coldlando@hotmail.com>
6
6
  License: Apache License 2.0
@@ -29,6 +29,7 @@ Requires-Dist: regex>=2023.12.25
29
29
  Requires-Dist: tokenizers>=0.15.2
30
30
  Requires-Dist: json-repair>=0.29.4
31
31
  Requires-Dist: robust-downloader>=0.0.2
32
+ Requires-Dist: bchlib>=2.1.3
32
33
  Provides-Extra: testing
33
34
  Requires-Dist: pytest>=7.4.4; extra == "testing"
34
35
  Requires-Dist: pytest-asyncio>=0.23.4; extra == "testing"
@@ -75,7 +76,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
75
76
  ### Usage 🖥️
76
77
 
77
78
  ```shell
78
- pip install novelai-python
79
+ pip -U install novelai-python
79
80
  ```
80
81
 
81
82
  More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground)
@@ -122,7 +123,8 @@ from dotenv import load_dotenv
122
123
  from pydantic import SecretStr
123
124
 
124
125
  from novelai_python import APIError, LoginCredential
125
- from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset
126
+ from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset, AdvanceLLMSetting
127
+ from novelai_python.sdk.ai.generate._enum import get_model_preset
126
128
 
127
129
  load_dotenv()
128
130
  username = os.getenv("NOVELAI_USER", None)
@@ -141,8 +143,17 @@ async def chat(prompt: str):
141
143
  agent = LLM.build(
142
144
  prompt=prompt,
143
145
  model=model,
144
- parameters=None # Auto Select or get from preset
146
+ # parameters=None, # Auto Select or get from preset
147
+ parameters=get_model_preset(TextLLMModel.ERATO).get_all_presets()[0].parameters, # Select from enum preset
148
+ advanced_setting=AdvanceLLMSetting(
149
+ min_length=1,
150
+ max_length=None, # Auto
151
+ )
145
152
  )
153
+ # NOTE:parameter > advanced_setting, which logic in generate/__init__.py
154
+ # If you not pass the parameter, it will use the default preset.
155
+ # So if you want to set the generation params, you should pass your own params.
156
+ # Only if you want to use some params not affect the generation, you can use advanced_setting.
146
157
  result = await agent.request(session=login_credential)
147
158
  except APIError as e:
148
159
  raise Exception(f"Error: {e.message}")
@@ -184,6 +195,14 @@ print(tokenizer_package.decode(tokenizer_package.encode(t_text)))
184
195
 
185
196
  ```
186
197
 
198
+ ## About Nsfw 🚫
199
+
200
+ You might need some solutions for identifying NSFW content and adding a mosaic to prevent operational mishaps.
201
+
202
+ https://dghs-imgutils.deepghs.org/main/api_doc/detect/nudenet.html
203
+
204
+ https://dghs-imgutils.deepghs.org/main/api_doc/operate/censor.html
205
+
187
206
  ## Acknowledgements 🙏
188
207
 
189
208
  [BackEnd](https://api.novelai.net/docs)
@@ -39,7 +39,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
39
39
  ### Usage 🖥️
40
40
 
41
41
  ```shell
42
- pip install novelai-python
42
+ pip -U install novelai-python
43
43
  ```
44
44
 
45
45
  More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground)
@@ -86,7 +86,8 @@ from dotenv import load_dotenv
86
86
  from pydantic import SecretStr
87
87
 
88
88
  from novelai_python import APIError, LoginCredential
89
- from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset
89
+ from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset, AdvanceLLMSetting
90
+ from novelai_python.sdk.ai.generate._enum import get_model_preset
90
91
 
91
92
  load_dotenv()
92
93
  username = os.getenv("NOVELAI_USER", None)
@@ -105,8 +106,17 @@ async def chat(prompt: str):
105
106
  agent = LLM.build(
106
107
  prompt=prompt,
107
108
  model=model,
108
- parameters=None # Auto Select or get from preset
109
+ # parameters=None, # Auto Select or get from preset
110
+ parameters=get_model_preset(TextLLMModel.ERATO).get_all_presets()[0].parameters, # Select from enum preset
111
+ advanced_setting=AdvanceLLMSetting(
112
+ min_length=1,
113
+ max_length=None, # Auto
114
+ )
109
115
  )
116
+ # NOTE:parameter > advanced_setting, which logic in generate/__init__.py
117
+ # If you not pass the parameter, it will use the default preset.
118
+ # So if you want to set the generation params, you should pass your own params.
119
+ # Only if you want to use some params not affect the generation, you can use advanced_setting.
110
120
  result = await agent.request(session=login_credential)
111
121
  except APIError as e:
112
122
  raise Exception(f"Error: {e.message}")
@@ -148,6 +158,14 @@ print(tokenizer_package.decode(tokenizer_package.encode(t_text)))
148
158
 
149
159
  ```
150
160
 
161
+ ## About Nsfw 🚫
162
+
163
+ You might need some solutions for identifying NSFW content and adding a mosaic to prevent operational mishaps.
164
+
165
+ https://dghs-imgutils.deepghs.org/main/api_doc/detect/nudenet.html
166
+
167
+ https://dghs-imgutils.deepghs.org/main/api_doc/operate/censor.html
168
+
151
169
  ## Acknowledgements 🙏
152
170
 
153
171
  [BackEnd](https://api.novelai.net/docs)
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "novelai-python"
3
- version = "0.5.0"
3
+ version = "0.5.2"
4
4
  description = "NovelAI Python Binding With Pydantic"
5
5
  authors = [
6
6
  { name = "sudoskys", email = "coldlando@hotmail.com" },
@@ -28,6 +28,7 @@ dependencies = [
28
28
  "tokenizers>=0.15.2",
29
29
  "json-repair>=0.29.4",
30
30
  "robust-downloader>=0.0.2",
31
+ "bchlib>=2.1.3",
31
32
  ]
32
33
  requires-python = ">=3.9"
33
34
  readme = "README.md"
@@ -1,11 +1,11 @@
1
1
  import math
2
2
  import random
3
- from typing import List, Optional, Union
3
+ from typing import List, Optional
4
4
 
5
5
  from pydantic import BaseModel
6
6
 
7
7
  from novelai_python.sdk.ai._const import map, initialN, initial_n, step, newN
8
- from novelai_python.sdk.ai._enum import Sampler, Model, ModelGroups, get_model_group, ModelTypeAlias
8
+ from novelai_python.sdk.ai._enum import Sampler, ModelGroups, get_model_group, ModelTypeAlias
9
9
 
10
10
 
11
11
  class Args(BaseModel):
@@ -83,9 +83,16 @@ class LLM(ApiBaseModel):
83
83
 
84
84
  @model_validator(mode="after")
85
85
  def normalize_model(self):
86
+ if self.model in [
87
+ TextLLMModel.NEO_2B, TextLLMModel.J_6B, TextLLMModel.J_6B_V3, TextLLMModel.J_6B_V4,
88
+ TextLLMModel.GENJI_JP_6B, TextLLMModel.GENJI_JP_6B_V2, TextLLMModel.GENJI_PYTHON_6B,
89
+ TextLLMModel.EUTERPE_V0, TextLLMModel.EUTERPE_V2, TextLLMModel.KRAKE_V1, TextLLMModel.KRAKE_V2,
90
+ TextLLMModel.CASSANDRA, TextLLMModel.COMMENT_BOT, TextLLMModel.INFILL, TextLLMModel.CLIO
91
+ ]:
92
+ self.endpoint = "https://api.novelai.net"
86
93
  tokenizer = NaiTokenizer(get_tokenizer_model(self.model))
87
94
  model_group = get_llm_group(self.model)
88
- total_tokens = len(tokenizer.encode(self.input))
95
+ total_tokens = tokenizer.total_tokens()
89
96
  if isinstance(self.input, str):
90
97
  prompt = tokenizer.encode(self.input)
91
98
  dtype = "uint32" if self.model in [TextLLMModel.ERATO] else "uint16"
@@ -238,10 +245,10 @@ class LLM(ApiBaseModel):
238
245
 
239
246
  valid_sequences = []
240
247
  for cell in logit_bias_group_exp:
241
- if not any(token < 0 or token >= total_tokens for token in cell.sequence):
248
+ if any(token < 0 or token >= total_tokens for token in cell.sequence):
242
249
  # 超出范围的 Token
243
- logger.debug(
244
- f"Bias {cell} contains tokens that are out of range and will be ignored."
250
+ logger.trace(
251
+ f"Bias [{cell}] contains tokens that are out of range and will be ignored."
245
252
  )
246
253
  else:
247
254
  # 将有效偏置组添加到列表
@@ -258,8 +265,8 @@ class LLM(ApiBaseModel):
258
265
  valid_bad_words_ids = []
259
266
  self.parameters.bad_words_ids = bad_words_ids
260
267
  for ban_word in self.parameters.bad_words_ids:
261
- if not any(token < 0 or token >= total_tokens for token in ban_word):
262
- logger.warning(
268
+ if any(token < 0 or token >= total_tokens for token in ban_word):
269
+ logger.trace(
263
270
  f"Bad word {ban_word} contains tokens that are out of range and will be ignored."
264
271
  )
265
272
  else:
@@ -277,6 +284,10 @@ class LLM(ApiBaseModel):
277
284
  self.advanced_setting.num_logprobs = self.logprobs_count
278
285
  if not self.advanced_setting.max_length:
279
286
  self.advanced_setting.max_length = 40
287
+ if self.parameters.repetition_penalty_range == 0:
288
+ self.parameters.repetition_penalty_range = None
289
+ if self.parameters.repetition_penalty_slope == 0:
290
+ self.parameters.repetition_penalty_slope = None
280
291
  return self
281
292
 
282
293
  @classmethod
@@ -294,7 +305,15 @@ class LLM(ApiBaseModel):
294
305
  **kwargs
295
306
  ) -> "LLM":
296
307
  """
297
- Generate Prompt
308
+ Generate
309
+ :param prompt: prompt, exp "Hello, World!"
310
+ :param model: TextLLMModelTypeAlias
311
+ :param advanced_setting: AdvanceLLMSetting, default None means auto-select
312
+ :param parameters: LLMGenerationParams, default None means auto-select, use get_model_preset
313
+ :param default_bias: bool, default True
314
+ :param logprobs_count: int, default None
315
+ :param bias_dinkus_asterism: bool, default False
316
+ :param prefix: str, default "vanilla"
298
317
  """
299
318
  assert isinstance(prompt, str), "Prompt must be a string"
300
319
  created = cls(
@@ -358,7 +377,7 @@ class LLM(ApiBaseModel):
358
377
  message = response_data.get("error") or response_data.get(
359
378
  "message") or f"Server error with status_code {response.status_code}"
360
379
  status_code = response_data.get("statusCode", response.status_code)
361
- if status_code == 200:
380
+ if status_code == 200 or status_code == 201:
362
381
  output = response_data.get("output", None)
363
382
  if not output:
364
383
  raise APIError(
@@ -148,30 +148,12 @@ def get_repetition_penalty_whitelist(group: TextTokenizerGroup) -> List[str]:
148
148
  return []
149
149
 
150
150
 
151
- def get_default_preset(model: TextLLMModel):
152
- if model == TextLLMModel.J_6B_V4:
153
- return J6bv4.get_all_presets()[0]
154
- elif model == TextLLMModel.NEO_2B:
155
- return Neo2b.get_all_presets()[0]
156
- elif model == TextLLMModel.GENJI_PYTHON_6B:
157
- return Genjipython6b.get_all_presets()[0]
158
- elif model == TextLLMModel.GENJI_JP_6B_V2:
159
- return Genjijp6bv2.get_all_presets()[0]
160
- elif model == TextLLMModel.EUTERPE_V2:
161
- return Euterpev2.get_all_presets()[0]
162
- elif model == TextLLMModel.KRAKE_V1:
163
- return Krakev1.get_all_presets()[0]
164
- elif model == TextLLMModel.KRAKE_V2:
165
- return Krakev2.get_all_presets()[0]
166
- elif model in COLORS_LLM:
167
- return Colorful.get_all_presets()[0]
168
- elif model == TextLLMModel.CLIO:
169
- return Clio.get_all_presets()[0]
170
- elif model == TextLLMModel.KAYRA:
171
- return Kayra.get_all_presets()[0]
172
- elif model == TextLLMModel.ERATO:
173
- return Erato.get_all_presets()[0]
174
- return Default.get_all_presets()[0]
151
+ def get_model_preset(model: TextLLMModel) -> "ModelPresetManager":
152
+ return PresetMapping.get(model, Default)
153
+
154
+
155
+ def get_default_preset(model: TextLLMModel) -> Preset:
156
+ return get_model_preset(model).get_all_presets()[0]
175
157
 
176
158
 
177
159
  class ModelPresetManager:
@@ -2279,3 +2261,24 @@ Erato.register_preset(
2279
2261
  description="Designed for writing in Japanese. Works fine for English too."
2280
2262
  )
2281
2263
  )
2264
+
2265
+ PresetMapping = {
2266
+ TextLLMModel.J_6B_V4: J6bv4,
2267
+ TextLLMModel.NEO_2B: Neo2b,
2268
+ TextLLMModel.GENJI_PYTHON_6B: Genjipython6b,
2269
+ TextLLMModel.GENJI_JP_6B_V2: Genjijp6bv2,
2270
+ TextLLMModel.EUTERPE_V2: Euterpev2,
2271
+ TextLLMModel.KRAKE_V1: Krakev1,
2272
+ TextLLMModel.KRAKE_V2: Krakev2,
2273
+ TextLLMModel.BLUE: Colorful,
2274
+ TextLLMModel.RED: Colorful,
2275
+ TextLLMModel.GREEN: Colorful,
2276
+ TextLLMModel.PURPLE: Colorful,
2277
+ TextLLMModel.PINK: Colorful,
2278
+ TextLLMModel.YELLOW: Colorful,
2279
+ TextLLMModel.WHITE: Colorful,
2280
+ TextLLMModel.BLACK: Colorful,
2281
+ TextLLMModel.CLIO: Clio,
2282
+ TextLLMModel.KAYRA: Kayra,
2283
+ TextLLMModel.ERATO: Erato,
2284
+ }
@@ -2,7 +2,7 @@ from dataclasses import dataclass
2
2
  from enum import Enum
3
3
  from typing import List, Optional, Tuple, Union
4
4
 
5
- from pydantic import BaseModel, Field, model_validator, ConfigDict
5
+ from pydantic import BaseModel, Field, ConfigDict
6
6
 
7
7
  """
8
8
  class StorySettings(BaseModel):
@@ -77,11 +77,6 @@ class LogitBiasGroup(BaseModel):
77
77
  ensure_sequence_finish: bool
78
78
  generate_once: bool
79
79
 
80
- @model_validator(mode="before")
81
- def validate_sequence(cls, value):
82
- print(value)
83
- return value
84
-
85
80
 
86
81
  class AdvanceLLMSetting(BaseModel):
87
82
  """
@@ -89,7 +84,7 @@ class AdvanceLLMSetting(BaseModel):
89
84
  """
90
85
  min_length: Optional[int] = 1
91
86
  max_length: Optional[int] = None
92
- repetition_penalty: Optional[float] = None
87
+ repetition_penalty: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
93
88
  generate_until_sentence: Optional[bool] = True
94
89
  use_cache: Optional[bool] = False
95
90
  use_string: Optional[bool] = False
@@ -107,36 +102,36 @@ class LLMGenerationParams(BaseModel):
107
102
  LLM Generation Settings
108
103
  """
109
104
  textGenerationSettingsVersion: Optional[int] = None
110
- temperature: Optional[float] = None
105
+ temperature: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
111
106
  max_length: Optional[int] = Field(default=None, ge=1, le=100) # Max 150(vip3),100(vip2)
112
107
  min_length: Optional[int] = Field(default=None, ge=1, le=100) # Max 150(vip3),100(vip2)
113
- top_k: Optional[float] = Field(default=None, ge=0)
114
- top_p: Optional[float] = Field(default=None, gt=0, le=1)
115
- top_a: Optional[float] = Field(default=None, ge=0)
116
- typical_p: Optional[float] = Field(default=None, ge=0, le=1)
117
- tail_free_sampling: Optional[float] = Field(default=None, ge=0)
118
- repetition_penalty: Optional[float] = Field(default=None, gt=0)
108
+ top_k: Optional[int] = Field(default=None, ge=0)
109
+ top_p: Optional[Union[float, int]] = Field(default=None, gt=0, le=1, allow_inf_nan=False)
110
+ top_a: Optional[Union[float, int]] = Field(default=None, ge=0, allow_inf_nan=False)
111
+ typical_p: Optional[Union[float, int]] = Field(default=None, ge=0, le=1, allow_inf_nan=False)
112
+ tail_free_sampling: Optional[Union[float, int]] = Field(default=None, ge=0, allow_inf_nan=False)
113
+ repetition_penalty: Optional[Union[float, int]] = Field(default=None, gt=0, allow_inf_nan=False)
119
114
  repetition_penalty_range: Optional[int] = Field(default=None, ge=0)
120
- repetition_penalty_slope: Optional[float] = Field(default=None, ge=0)
115
+ repetition_penalty_slope: Optional[Union[float, int]] = Field(default=None, ge=0, allow_inf_nan=False)
121
116
 
122
117
  eos_token_id: int = None
123
118
  bad_words_ids: List[List[int]] = None
124
119
  logit_bias_groups: Optional[List[LogitBiasGroup]] = []
125
120
 
126
- repetition_penalty_frequency: Optional[float] = None
127
- repetition_penalty_presence: Optional[float] = None
121
+ repetition_penalty_frequency: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
122
+ repetition_penalty_presence: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
128
123
  repetition_penalty_whitelist: Optional[List[int]] = None
129
124
  repetition_penalty_default_whitelist: Optional[bool] = None
130
- cfg_scale: Optional[float] = None
125
+ cfg_scale: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
131
126
  cfg_uc: Optional[str] = None
132
127
  phrase_rep_pen: PenStyle = PenStyle.Off
133
- top_g: Optional[float] = None
134
- mirostat_tau: Optional[float] = None
135
- mirostat_lr: Optional[float] = None
136
- math1_temp: Optional[float] = None
137
- math1_quad: Optional[float] = None
138
- math1_quad_entropy_scale: Optional[float] = None
139
- min_p: Optional[float] = None
128
+ top_g: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
129
+ mirostat_tau: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
130
+ mirostat_lr: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
131
+ math1_temp: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
132
+ math1_quad: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
133
+ math1_quad_entropy_scale: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
134
+ min_p: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
140
135
 
141
136
  order: Union[List[int], List[KeyOrderEntry]] = [
142
137
  KeyOrderEntry(id=Key.Cfg, enabled=False),
@@ -134,7 +134,7 @@ class LLMStream(LLM):
134
134
  request=request_data, code=status_code, response=message
135
135
  )
136
136
  else:
137
- if response.status_code != 200:
137
+ if response.status_code not in [200, 201]:
138
138
  raise APIError(
139
139
  f"Server error with status code {response.status_code}",
140
140
  request=request_data, code=response.status_code, response=response.content
@@ -187,6 +187,13 @@ class NaiTokenizer:
187
187
  return self.tokenizer.encode(text).tokens
188
188
  raise NotImplementedError("Tokenizer does not support token encoding")
189
189
 
190
+ def total_tokens(self) -> int:
191
+ if isinstance(self.tokenizer, Tokenizer):
192
+ return len(self.tokenizer.get_vocab())
193
+ if isinstance(self.tokenizer, SimpleTokenizer):
194
+ return len(self.tokenizer.get_vocab())
195
+ raise NotImplementedError("Tokenizer does not support token encoding")
196
+
190
197
  def encode(self, sentence: str) -> List[int]:
191
198
  if isinstance(self.tokenizer, SimpleTokenizer):
192
199
  return self.tokenizer.encode(sentence).ids
@@ -144,3 +144,11 @@ class SimpleTokenizer(object):
144
144
  text = ''.join([self.decoder[token] for token in tokens])
145
145
  text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
146
146
  return text
147
+
148
+ def get_vocab(self):
149
+ return self.encoder
150
+
151
+
152
+ if __name__ == '__main__':
153
+ tokenizer = SimpleTokenizer()
154
+ print(len(tokenizer.get_vocab()))