From 8ca0fa157976fe6ad34cc8ed7707f205b7bda573 Mon Sep 17 00:00:00 2001 From: "chamaael@etu.unistra.fr" Date: Sat, 16 Mar 2024 17:34:02 +0100 Subject: [PATCH] Flake8 corrections --- sdk/models/model_text_conversation.py | 12 ++++++------ sdk/options/options_text_conversation.py | 3 +-- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/sdk/models/model_text_conversation.py b/sdk/models/model_text_conversation.py index 3b39df7..b54c302 100644 --- a/sdk/models/model_text_conversation.py +++ b/sdk/models/model_text_conversation.py @@ -1,7 +1,6 @@ import torch -from typing import Optional, Dict, Tuple -from transformers import Conversation, AutoModelForCausalLM -from sdk.options.options_tokenizer import OptionsTokenizer +from typing import Optional, Dict +from transformers import AutoModelForCausalLM from sdk.tokenizers.tokenizer import Tokenizer from sdk.models import Model from sdk.options import Devices, OptionsTextConversation @@ -52,7 +51,8 @@ class ModelsTextConversation(Model): conversation_active: bool = False - def __init__(self, model_name: str, model_path: str, option: OptionsTextConversation): + def __init__(self, model_name: str, model_path: str, + option: OptionsTextConversation): """ Initializes the ModelsTextToImage class :param model_name: The name of the model @@ -72,8 +72,8 @@ def create_pipeline(self): self.pipeline = AutoModelForCausalLM.from_pretrained( self.model_path, trust_remote_code=self.options.trust_remote_code, - pad_token_id=self.options.pad_token_id, # Custom pad token ID - eos_token_id=self.options.eos_token_id # Custom end-of-sequence token ID + pad_token_id=self.options.pad_token_id, + eos_token_id=self.options.eos_token_id ) def load_model(self, option: OptionsTextConversation) -> bool: diff --git a/sdk/options/options_text_conversation.py b/sdk/options/options_text_conversation.py index 8a7bd83..73828c3 100644 --- a/sdk/options/options_text_conversation.py +++ b/sdk/options/options_text_conversation.py @@ -104,5 +104,4 @@ def __init__(self, self.delete_conv = delete_conv self.tokenizer_id_to_use = tokenizer_id_to_use self.chat_id_to_use = chat_id_to_use - self.trust_remote_code = trust_remote_code - + self.trust_remote_code = trust_remote_code \ No newline at end of file