[AI] stable-code-3b 기본적인 사용 가이드 (AI coding)

[AI] stable-code-3b 기본적인 사용 가이드 (AI coding)

https://huggingface.co/stabilityai/stable-code-3b

stabilityai/stable-code-3b · Hugging Face

아래 링크로 colab 접속해서 기본 가이드 작성 해보자

https://colab.research.google.com/?hl=ko

Google Colaboratory

의존 패키지 설치

!pip install -q transformers accelerate sentencepiece

Load 모델 

import psutil
from transformers import AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline

model = AutoModelForCausalLM.from_pretrained(
    "stabilityai/stable-code-3b",
    trust_remote_code=True,
    torch_dtype="auto",
)
tokenizer = AutoTokenizer.from_pretrained(
    "stabilityai/stable-code-3b",
    trust_remote_code=True
)
pipeline = TextGenerationPipeline(
    model=model,
    tokenizer=tokenizer,
    device=0,  # gpu
    # device=-1,  # cpu
)


def generate(instruction, max_token=2000):
    channel = pipeline(
        instruction,
        num_workers=psutil.cpu_count(),
        do_sample=True,
        top_k=10,
        temperature=0.2,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
        return_full_text=True,
        truncation=True,
        max_length=max_token,
    )

    for ch in channel:
        code = ch['generated_text'].strip()
        print("\n".join(code.split('\n')), flush=True)

원하는 텍스트 입력

generate("Please write a binary search algorithm function using python")

실행 결과