taide

/

Llama3-TAIDE-LX-8B-Chat-Alpha1

like 119

Follow

TAIDE 184

本地機、colab(T4-GPU)都嘗試執行過以下程式，但是都運行不出結果(不斷運行、沒有任何錯誤異常)，不知道是否哪邊出錯?

import torch
from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

login(
    token="HF_TOKEN",
    add_to_git_credential=True
)

model_id = "taide/Llama3-TAIDE-LX-8B-Chat-Alpha1"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
)



@torch

	.inference_mode()
def get_completion(prompt: str) -> str:
    response = pipe(
        prompt,
        max_new_tokens=256,
        do_sample=True, 
        temperature=0.1, 
        top_k=50, 
        top_p=0.1, 
        eos_token_id=pipe.tokenizer.eos_token_id, 
        pad_token_id=pipe.tokenizer.pad_token_id
    )
    return response

response = get_completion("你好")
print(response)