diff --git a/backend/oasst_backend/api/v1/hugging_face.py b/backend/oasst_backend/api/v1/hugging_face.py index a8d8aeb9..62d2ea6b 100644 --- a/backend/oasst_backend/api/v1/hugging_face.py +++ b/backend/oasst_backend/api/v1/hugging_face.py @@ -4,7 +4,7 @@ from fastapi import APIRouter, Depends from oasst_backend.api import deps from oasst_backend.models import ApiClient from oasst_backend.schemas.hugging_face import ToxicityClassification -from oasst_backend.utils.hugging_face import HF_url, HuggingFaceAPI +from oasst_backend.utils.hugging_face import HfUrl, HuggingFaceAPI router = APIRouter() @@ -25,7 +25,7 @@ async def get_text_toxicity( ToxicityClassification: the score of toxicity of the message. """ - api_url: str = HF_url.HUGGINGFACE_TOXIC_ROBERTA.value + api_url: str = HfUrl.HUGGINGFACE_TOXIC_ROBERTA.value hugging_face_api = HuggingFaceAPI(api_url) response = await hugging_face_api.post(msg) diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py index 186a6ac8..821ba562 100644 --- a/backend/oasst_backend/api/v1/tasks.py +++ b/backend/oasst_backend/api/v1/tasks.py @@ -9,7 +9,7 @@ from oasst_backend.api import deps from oasst_backend.api.v1.utils import prepare_conversation from oasst_backend.config import settings from oasst_backend.prompt_repository import PromptRepository, TaskRepository -from oasst_backend.utils.hugging_face import HF_embeddingModel, HF_url, HuggingFaceAPI +from oasst_backend.utils.hugging_face import HfEmbeddingModel, HfUrl, HuggingFaceAPI from oasst_shared.exceptions import OasstError, OasstErrorCode from oasst_shared.schemas import protocol as protocol_schema from sqlmodel import Session @@ -285,15 +285,15 @@ async def tasks_interaction( if not settings.DEBUG_SKIP_EMBEDDING_COMPUTATION: try: hugging_face_api = HuggingFaceAPI( - f"{HF_url.HUGGINGFACE_FEATURE_EXTRACTION.value}/{HF_embeddingModel.MINILM.value}" + f"{HfUrl.HUGGINGFACE_FEATURE_EXTRACTION.value}/{HfEmbeddingModel.MINILM.value}" ) embedding = await hugging_face_api.post(interaction.text) pr.insert_message_embedding( - message_id=newMessage.id, model=HF_embeddingModel.MINILM.value, embedding=embedding + message_id=newMessage.id, model=HfEmbeddingModel.MINILM.value, embedding=embedding ) except OasstError: logger.error( - f"Could not fetch embbeddings for text reply to {interaction.message_id=} with {interaction.text=} by {interaction.user=}." + f"Could not fetch embbeddings for text reply to {interaction.message_id=} with {interaction.text=} by {interaction.user=}." ) return protocol_schema.TaskDone() diff --git a/backend/oasst_backend/utils/hugging_face.py b/backend/oasst_backend/utils/hugging_face.py index ef440055..87c6288e 100644 --- a/backend/oasst_backend/utils/hugging_face.py +++ b/backend/oasst_backend/utils/hugging_face.py @@ -7,12 +7,12 @@ from oasst_backend.config import settings from oasst_shared.exceptions import OasstError, OasstErrorCode -class HF_url(str, Enum): +class HfUrl(str, Enum): HUGGINGFACE_TOXIC_ROBERTA = ("https://api-inference.huggingface.co/models/unitary/multilingual-toxic-xlm-roberta",) HUGGINGFACE_FEATURE_EXTRACTION = "https://api-inference.huggingface.co/pipeline/feature-extraction" -class HF_embeddingModel(str, Enum): +class HfEmbeddingModel(str, Enum): MINILM = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"