mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-15 00:07:28 +00:00
Model: Add support for HuggingFace config and bad_words_ids
This is necessary for Kobold's API. Current models use bad_words_ids in generation_config.json, but for some reason, they're also present in the model's config.json. Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
import json
|
||||
import pathlib
|
||||
from typing import List, Optional, Union
|
||||
from loguru import logger
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
@@ -11,6 +12,7 @@ class GenerationConfig(BaseModel):
|
||||
"""
|
||||
|
||||
eos_token_id: Optional[Union[int, List[int]]] = None
|
||||
bad_words_ids: Optional[List[List[int]]] = None
|
||||
|
||||
@classmethod
|
||||
def from_file(self, model_directory: pathlib.Path):
|
||||
@@ -30,3 +32,40 @@ class GenerationConfig(BaseModel):
|
||||
return [self.eos_token_id]
|
||||
else:
|
||||
return self.eos_token_id
|
||||
|
||||
|
||||
class HuggingFaceConfig(BaseModel):
|
||||
"""
|
||||
An abridged version of HuggingFace's model config.
|
||||
Will be expanded as needed.
|
||||
"""
|
||||
|
||||
badwordsids: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def from_file(self, model_directory: pathlib.Path):
|
||||
"""Create an instance from a generation config file."""
|
||||
|
||||
hf_config_path = model_directory / "config.json"
|
||||
with open(
|
||||
hf_config_path, "r", encoding="utf8"
|
||||
) as hf_config_json:
|
||||
hf_config_dict = json.load(hf_config_json)
|
||||
return self.model_validate(hf_config_dict)
|
||||
|
||||
def get_badwordsids(self):
|
||||
"""Wrapper method to fetch badwordsids."""
|
||||
|
||||
if self.badwordsids:
|
||||
try:
|
||||
bad_words_list = json.loads(self.badwordsids)
|
||||
return bad_words_list
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(
|
||||
"Skipping badwordsids from config.json "
|
||||
"since it's not a valid array."
|
||||
)
|
||||
|
||||
return []
|
||||
else:
|
||||
return []
|
||||
|
||||
Reference in New Issue
Block a user