Refactor the sampling class (#199)

* improve validation

* remove to_gen_params functions

* update changes for all endpoint types

* OAI: Fix calls to generation

Chat completion and completion need to have prompt split out before
pushing to the backend.

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Convert Top-K values of -1 to 0

Some OAI implementations use -1 as disabled instead of 0. Therefore,
add a coalesce case.

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Format and space out

Make the code more readable.

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Fix mirostat

Field items are nested in data within a Pydantic FieldInfo

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Format

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Fix banned_tokens and allowed_tokens conversion

If the provided string has whitespace, trim it before splitting.

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Add helpful log to dry_sequence_breakers

Let the user know if the sequence errors out.

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Apply validators in right order

Validators need to be applied in order from top to bottom, this is why
the after validator was not being applied properly.

Set the model to validate default params for sampler override purposes.
This can be turned off if there are unclear errors.

Signed-off-by: kingbri <bdashore3@proton.me>

* Endpoints: Format

Cleanup and semantically fix field validators

Signed-off-by: kingbri <bdashore3@proton.me>

* Kobold: Update validators and fix parameter application

Validators on parent fields cannot see child fields. Therefore,
validate using the child fields instead and alter the parent field
data from there.

Also fix badwordsids casting.

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Remove validate defaults and fix mirostat

If a user sets an override to a non-default value, that's their
own fault.

Run validator on the actual mirostat_mode parameter rather than
the alternate mirostat parameter.

Signed-off-by: kingbri <bdashore3@proton.me>

* Kobold: Rework badwordsids

Currently, this serves to ban the EOS token. All other functionality
was legacy, so remove it.

Signed-off-by: kingbri <bdashore3@proton.me>

* Model: Remove HuggingfaceConfig

This was only necessary for badwordsids. All other fields are handled
by exl2. Keep the class as a stub if it's needed again.

Signed-off-by: kingbri <bdashore3@proton.me>

* Kobold: Bump kcpp impersonation

TabbyAPI supports XTC now.

Signed-off-by: kingbri <bdashore3@proton.me>

* Sampling: Change alias to validation_alias

Reduces the probability for errors and makes the class consistent.

Signed-off-by: kingbri <bdashore3@proton.me>

* OAI: Use constraints for validation

Instead of adding a model_validator, use greater than or equal to
constraints provided by Pydantic.

Signed-off-by: kingbri <bdashore3@proton.me>

* Tree: Lint

Signed-off-by: kingbri <bdashore3@proton.me>

---------

Co-authored-by: SecretiveShell <84923604+SecretiveShell@users.noreply.github.com>
Co-authored-by: kingbri <bdashore3@proton.me>

endpoints/Kobold/router.py

@@ -137,7 +137,7 @@ async def get_version():
 async def get_extra_version():
     """Impersonate Koboldcpp."""
 
-    return {"result": "KoboldCpp", "version": "1.71"}
+    return {"result": "KoboldCpp", "version": "1.74"}
 
 
 @kai_router.get("/config/soft_prompts_list")

endpoints/Kobold/types/generation.py

@@ -1,9 +1,9 @@
+from functools import partial
+from pydantic import BaseModel, Field, field_validator
 from typing import List, Optional
 
-from pydantic import BaseModel, Field
-from common import model
 from common.sampling import BaseSamplerRequest, get_default_sampler_value
-from common.utils import flat_map, unwrap
+from common.utils import unwrap
 
 
 class GenerateRequest(BaseSamplerRequest):
@@ -11,29 +11,31 @@ class GenerateRequest(BaseSamplerRequest):
     genkey: Optional[str] = None
     use_default_badwordsids: Optional[bool] = False
     dynatemp_range: Optional[float] = Field(
-        default_factory=get_default_sampler_value("dynatemp_range")
+        default_factory=partial(get_default_sampler_value, "dynatemp_range")
     )
 
-    def to_gen_params(self, **kwargs):
-        # Exl2 uses -1 to include all tokens in repetition penalty
-        if self.penalty_range == 0:
-            self.penalty_range = -1
+    # Validate on the parent class's fields
+    @field_validator("penalty_range", mode="before")
+    def validate_penalty_range(cls, v):
+        return -1 if v == 0 else v
 
-        if self.dynatemp_range:
-            self.min_temp = self.temperature - self.dynatemp_range
-            self.max_temp = self.temperature + self.dynatemp_range
+    @field_validator("dynatemp_range", mode="before")
+    def validate_temp_range(cls, v, field_info):
+        if v > 0:
+            # A default temperature is always 1
+            temperature = unwrap(field_info.data.get("temperature"), 1)
 
-        # Move badwordsids into banned tokens for generation
-        if self.use_default_badwordsids:
-            bad_words_ids = unwrap(
-                model.container.generation_config.bad_words_ids,
-                model.container.hf_config.get_badwordsids(),
-            )
+            field_info.data["min_temp"] = temperature - v
+            field_info.data["max_temp"] = temperature + v
 
-            if bad_words_ids:
-                self.banned_tokens += flat_map(bad_words_ids)
+        return v
 
-        return super().to_gen_params(**kwargs)
+    # Currently only serves to ban EOS token, but can change
+    @field_validator("use_default_badwordsids", mode="before")
+    def validate_badwordsids(cls, v, field_info):
+        field_info.data["ban_eos_token"] = v
+
+        return v
 
 
 class GenerateResponseResult(BaseModel):

endpoints/Kobold/utils/generation.py

@@ -53,7 +53,7 @@ async def _stream_collector(data: GenerateRequest, request: Request):
         logger.info(f"Received Kobold generation request {data.genkey}")
 
         generator = model.container.generate_gen(
-            data.prompt, data.genkey, abort_event, **data.to_gen_params()
+            request_id=data.genkey, abort_event=abort_event, **data.model_dump()
         )
         async for generation in generator:
             if disconnect_task.done():