import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

model_id = "google/gemma-7b"
peft_model_id = "brildev7/gemma-7b-translation-enko-sft-qlora"
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=False
)

model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    quantization_config=quantization_config, 
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
    token=os.environ['HF_TOKEN'],
    device_map="auto"
)
model = PeftModel.from_pretrained(model, peft_model_id)

tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
tokenizer.pad_token_id = tokenizer.eos_token_id

# example
prompt_template = """Translate the following sentences into Korean language:
{}

translation:
"""
sentences = "Apple is facing a crisis in one of its key markets, China, as it is being challenged by local smartphone manufacturers. In a bid to counter the threat, Apple CEO Tim Cook is reportedly planning to visit China to meet with local smartphone manufacturers and discuss a joint investment. Apple is also reportedly considering installing an AI model from Baidu, the Chinese search giant, on its iPhone. The move comes as Apple is facing a price war in China, with local smartphone manufacturers offering steep discounts on their products."
texts = prompt_template.format(sentences)
inputs = tokenizer(texts, return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, max_new_tokens=1024)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
- 애플은 국산 스마트폰 제조사들의 도발에 중국에서 하나의 핵심 시장에 위기를 맞고 있다. 이 위협을 타개하기 위해 애플의 최고 경영자인 팀 쿡은 중국을 방문해 현지 스마트폰 제조사들과 접촉해 공동 투자를 논의하는 것으로 알려졌다. 애플은 또한 중국 최대 검색사 바이두(Baidu)의 인공 지능(AI) 모델을 아이폰에 탑재하는 것을 검토 중인 것으로 전해졌다. 애플은 국내 스마트폰 제조사들이 자신들의 제품에 급한 할인을 내놓으면서 중국에서 가격전쟁에 직면해 있는 것이다.

# example
sentences = "Is it safe to drink milk and eat chicken?"
texts = prompt_template.format(sentences)
inputs = tokenizer(texts, return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, max_new_tokens=1024)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
- 우유와 닭고기는 안전한가요?

# example
sentences = "What precautions to take during the bird flu outbreak"
texts = prompt_template.format(sentences)
inputs = tokenizer(texts, return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, max_new_tokens=1024)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
- 조류 독감 유행 시 어떠한 주의 사항을 해야 하는지