Views
No views yet
1from transformers import pipeline, AutoModelForCausalLM
2import torch
3
4LLM_MODEL = "NTIS/KoRnDAlpaca-Polyglot-12.8B"
5query = "지능형 영상감시 기술의 대표적인 국내 기업은?"
6
7llm_model = AutoModelForCausalLM.from_pretrained(
8 LLM_MODEL,
9 device_map="auto",
10 torch_dtype=torch.float16,
11 low_cpu_mem_usage=True,
12 # load_in_8bit=True,
13 # revision="8bit"
14)
15pipe = pipeline(
16 "text-generation",
17 model=llm_model,
18 tokenizer=LLM_MODEL,
19 # device=2,
20)
21
22ans = pipe(
23 f"### 질문: {query}\n\n### 답변:",
24 do_sample=True,
25 max_new_tokens=512,
26 temperature=0.1,
27 top_p=0.9,
28 return_full_text=False,
29 eos_token_id=2,
30 )
31msg = ans[0]["generated_text"]
32
33if len(msg.split('###')[0]) > 0:
34 output = msg.split('###')[0]
35else:
36 output = '답변을 드리지 못하여 죄송합니다.'
37
38print(output)
39# 국내 지능형 영상감시 기술의 대표적인 기업으로는 한화 테크윈이 있다.