mintaeng commited on
Commit
651da14
1 Parent(s): d09dce6

Create load_model.py

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