Views
No views yet
### 한국어: {sentence}</끝>
### 영어:### 영어: {sentence}</끝>
### 한국어:1from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
2import torch
3repo = "squarelike/Gugugo-koen-7B-V1.1-GPTQ"
4model = AutoModelForCausalLM.from_pretrained(
5 repo,
6 device_map='auto'
7)
8tokenizer = AutoTokenizer.from_pretrained(repo)
9
10model.eval()
11model.config.use_cache = True
12
13class StoppingCriteriaSub(StoppingCriteria):
14 def __init__(self, stops = [], encounters=1):
15 super().__init__()
16 self.stops = [stop for stop in stops]
17
18 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
19 for stop in self.stops:
20 if torch.all((stop == input_ids[0][-len(stop):])).item():
21 return True
22
23 return False
24
25stop_words_ids = torch.tensor([[829, 45107, 29958], [1533, 45107, 29958], [829, 45107, 29958], [21106, 45107, 29958]]).to("cuda")
26stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
27
28def gen(lan="en", x=""):
29 if (lan == "ko"):
30 prompt = f"### 한국어: {x}</끝>\n### 영어:"
31 else:
32 prompt = f"### 영어: {x}</끝>\n### 한국어:"
33 gened = model.generate(
34 **tokenizer(
35 prompt,
36 return_tensors='pt',
37 return_token_type_ids=False
38 ).to("cuda"),
39 max_new_tokens=2000,
40 temperature=0.3,
41 # no_repeat_ngram_size=5,
42 num_beams=5,
43 stopping_criteria=stopping_criteria
44 )
45 return tokenizer.decode(gened[0][1:]).replace(prompt+" ", "").replace("</끝>", "")
46
47
48print(gen(lan="en", x="Hello, world!"))