This commit is contained in:
itqop 2025-03-03 04:11:07 +03:00
parent 115aea4792
commit 6a72b0404b
2 changed files with 3 additions and 4 deletions

View File

@ -11,5 +11,5 @@ TEXT_TASK_CHANNEL = os.environ.get("TEXT_TASK_CHANNEL", "text_task_channel")
TEXT_RESULT_CHANNEL = os.environ.get("TEXT_RESULT_CHANNEL", "text_result_channel") TEXT_RESULT_CHANNEL = os.environ.get("TEXT_RESULT_CHANNEL", "text_result_channel")
BASE_MODEL = os.environ.get("BASE_MODEL", "google/gemma-2-2b") BASE_MODEL = os.environ.get("BASE_MODEL", "google/gemma-2-2b")
ADAPTER_DIR = os.environ.get("ADAPTER_DIR", "./gemma-2-2b_lora") ADAPTER_DIR = os.environ.get("ADAPTER_DIR", "itqop/gemma-2-2b-summarize-adapter")
HF_TOKEN = os.environ.get("HF_TOKEN", "") HF_TOKEN = os.environ.get("HF_TOKEN", "")

View File

@ -32,14 +32,13 @@ class ModelLoader:
self.model = PeftModel.from_pretrained( self.model = PeftModel.from_pretrained(
base_model, base_model,
self.adapter_dir, self.adapter_dir
local_files_only=True
) )
self.model.eval() self.model.eval()
self.model.to(self.device) self.model.to(self.device)
def load_tokenizer(self): def load_tokenizer(self):
self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_name, local_files_only=True) self.tokenizer = AutoTokenizer.from_pretrained(self.adapter_dir)
if self.tokenizer.pad_token is None: if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token self.tokenizer.pad_token = self.tokenizer.eos_token