mintaeng commited on
Commit
c04e74d
1 Parent(s): c5c9e79

Create load_model.py

Browse files
Files changed (1) hide show
  1. load_model.py +10 -19
load_model.py CHANGED
@@ -2,28 +2,19 @@ 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(
 
2
  import torch
3
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
4
  from langchain_huggingface.llms import HuggingFacePipeline
 
 
 
 
 
5
 
6
  '''
7
+ AutoModelForCausalLm을 사용하여 모델 불러오기
 
 
 
 
 
 
 
8
 
9
+ 속도가 느린편에 속하나 문제 없이 돌아갈 수 있는 안정적인 방식
10
+ '''
11
+ def load_Auto():
12
+ model_id = 'Dongwookss/small_fut_final'
13
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
14
+ model = AutoModelForCausalLM.from_pretrained(
15
+ model_id,
16
+ torch_dtype=torch.bfloat16,
17
+ device_map="auto",
18
  )
19
 
20
  pipe = pipeline(