airllm icon indicating copy to clipboard operation
airllm copied to clipboard

How to set system prompt

Open OKHand-Zy opened this issue 5 months ago • 1 comments

I tried setting up a system prompt, but the generated results were unsatisfactory. I'm wondering if there's something missing in my code settings.

this is my code :

from airllm import AutoModel
import mlx.core as mx
import time

#MAX_LENGTH = 128 
MAX_NEW_TOKENS = 5120 

# could use hugging face model repo id:
model = AutoModel.from_pretrained(f"/Users/ziyu/RemoteFolder/LLM/LLM-Model/Ori/Meta-Llama-3.1-8B-Instruct/",
    layer_shards_saving_path=r"/Users/ziyu/RemoteFolder/LLM/airLLM/")

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"

DEFAULT_SYSTEM_PROMPT = """請詳細的描述與分析這個故事的內容,並使用「繁體中文」回答"""

input_prompt=""" ... """

def tokenize(model, text):
    input_ids = model.tokenizer(
        text,
        return_tensors="np", 
        return_attention_mask=False,
        truncation=True,
        #max_length=MAX_LENGTH,
        #padding=False
    )
    return input_ids

def generate(model, input_ids):
    generation_output = model.generate(
        mx.array(input_ids['input_ids']),
        max_new_tokens=MAX_NEW_TOKENS,
        use_cache=True,
        return_dict_in_generate=True,
        temperature=0.0,
        top_p=0.9,
        do_sample=True,
        num_beams=1,
        repetition_penalty=1.2,
    )
    return generation_output

def query(model, text):
    input_ids = tokenize(model, text)
    generation_output = generate(model, input_ids)
    return generation_output

messages = "{b_inst} {system}{prompt} {e_inst}".format(
        b_inst=B_INST,
        system=f"{B_SYS}{DEFAULT_SYSTEM_PROMPT}{E_SYS}",
        prompt=input_prompt,
        e_inst=E_INST)

start = time.process_time()
output = query(model, messages)
end = time.process_time()
print('----------------')
print(output)
print('--- ', end - start, ' sec ---')

OKHand-Zy avatar Sep 02 '24 09:09 OKHand-Zy