Views
No views yet
AutoModelForCausalLM and AutoTokenizer1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3# Load the model and tokenizer
4tokenizer = AutoTokenizer.from_pretrained("prem-research/prem-1B-chat")
5model = AutoModelForCausalLM.from_pretrained('prem-research/prem-1B-chat', torch_dtype=torch.bfloat16)
6model = model.to('cuda')
7
8# Setup terminators
9terminators = [tokenizer.eos_token_id, tokenizer.encode('<|eot_id|>', add_special_tokens=False)[0]]
10
11# Prepare the prompt
12messages = [
13 {
14 "role": "system",
15 "content": "You are a helpful AI assistant. You should give concise responses to very simple questions, but provide thorough responses to more complex and open-ended questions."
16 },
17 {
18 'role': 'user',
19 'content': 'Help me understand machine learning.'
20 }
21]
22
23prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
24
25# Generate
26inputs = tokenizer(prompt, return_attention_mask=False, return_tensors="pt", add_special_tokens=False)
27input_ids = inputs['input_ids']
28input_ids = input_ids.to(model.device)
29res = model.generate(input_ids=input_ids, max_new_tokens=400, pad_token_id=tokenizer.pad_token_id, eos_token_id=terminators)
30generated_text = tokenizer.decode(res[0][input_ids.shape[1]:], skip_special_tokens=True).strip()
31print(generated_text)1import torch
2from transformers import pipeline
3
4# Load the pipeline
5pipe = pipeline("text-generation", model="prem-research/prem-1B-chat", torch_dtype=torch.bfloat16, device=0)
6
7# Prepare prompt
8messages = [
9 {
10 "role": "system",
11 "content": "You are a helpful AI assistant. You should give concise responses to very simple questions, but provide thorough responses to more complex and open-ended questions."
12 },
13 {
14 'role': 'user',
15 'content': 'Help me understand machine learning.'
16 }
17]
18prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
19
20# Setup terminators
21terminators = [pipe.tokenizer.eos_token_id, pipe.tokenizer.encode('<|eot_id|>', add_special_tokens=False)[0]]
22
23# Generate
24outputs = pipe(prompt, max_new_tokens=400, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, pad_token_id=pipe.tokenizer.pad_token_id, eos_token_id=terminators)
25print(outputs[0]["generated_text"][len(prompt):])| Model | Avg | Arc-c | Arc-e | Hellaswag | MMLU | Obqa | Piqa | Winogrande |
|---|---|---|---|---|---|---|---|---|
| prem-1B | 42.64 | 24.74 | 57.40 | 42.01 | 24.75 | 21.00 | 72.14 | 56.43 |
| prem-1B-chat | 41.76 | 24.48 | 53.32 | 40.28 | 25.27 | 22.20 | 70.89 | 55.88 |
| TinyLlama-1.1B-Chat-v1.0 | 46.16 | 30.03 | 61.53 | 46.56 | 24.72 | 25.80 | 74.21 | 60.29 |
| opt-1.3b | 42.94 | 23.37 | 57.44 | 41.49 | 24.86 | 23.20 | 71.49 | 58.72 |
| pythia-1b | 40.71 | 24.31 | 56.90 | 37.72 | 23.20 | 18.80 | 70.62 | 53.43 |
