mintaeng commited on
Commit
cde4770
1 Parent(s): d9df740

Create load_model.py

Browse files
Files changed (1) hide show
  1. load_model.py +40 -0
load_model.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
4
+ from langchain_huggingface.llms import HuggingFacePipeline
5
+ from unsloth import FastLanguageModel
6
+
7
+ import dotenv
8
+ dotenv()
9
+
10
+
11
+ '''
12
+ FastLanguageModel 사용하여 모델 불러오기
13
+ Fine-tuning을 도와주는 Unsloth 패키지를 사용하여 모델 불러오기
14
+ 불러오는 모델 repo안에 adapter_config가 존재하면 안된다.
15
+ 빠른 추론 할 수 있도록 도와준다.
16
+ '''
17
+ def load_Fast():
18
+ max_seq_length = 2048
19
+ hf_token = os.getenv('hf_token')
20
+
21
+ model, tokenizer = FastLanguageModel.from_pretrained(
22
+ model_name="Dongwookss/last_small_pre", # adapter_config가 존재하지 않는 모델명으로 불러와야한다.
23
+ max_seq_length=max_seq_length,
24
+ dtype=None,
25
+ load_in_4bit=False,
26
+ token = hf_token,
27
+ )
28
+
29
+ pipe = pipeline(
30
+ "text-generation",
31
+ model=model,
32
+ tokenizer=tokenizer,
33
+ max_new_tokens=2048,
34
+ )
35
+
36
+ llm = HuggingFacePipeline(
37
+ pipeline=pipe,
38
+ model_kwargs={"temperature": 0.7},
39
+ )
40
+ return llm