Views
No views yet
1import torch
2
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5from conversation import get_default_conv_template
6
7# MiniChat
8tokenizer = AutoTokenizer.from_pretrained("GeneZC/MiniChat-3B", use_fast=False)
9# GPU.
10model = AutoModelForCausalLM.from_pretrained("GeneZC/MiniChat-3B", use_cache=True, device_map="auto", torch_dtype=torch.float16).eval()
11# CPU.
12# model = AutoModelForCausalLM.from_pretrained("GeneZC/MiniChat-3B", use_cache=True, device_map="cpu", torch_dtype=torch.float32).eval()
13device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
14conv = get_default_conv_template("minichat")
15
16question = "Implement a program to find the common elements in two arrays without using any extra data structures."
17conv.append_message(conv.roles[0], question)
18conv.append_message(conv.roles[1], None)
19prompt = conv.get_prompt()
20input_ids = tokenizer([prompt]).input_ids
21output_ids = model.generate(
22 torch.as_tensor(input_ids).to(device),
23 do_sample=True,
24 temperature=0.7,
25 max_new_tokens=1024,
26)
27output_ids = output_ids[0][len(input_ids[0]):]
28output = tokenizer.decode(output_ids, skip_special_tokens=True).strip()
29# output: "def common_elements(arr1, arr2):\n if len(arr1) == 0:\n return []\n if len(arr2) == 0:\n return arr1\n\n common_elements = []\n for element in arr1:\n if element in arr2:\n common_elements.append(element)\n\n return common_elements"
30# Multiturn conversation could be realized by continuously appending questions to `conv`.1@article{zhang2023law,
2 title={Towards the Law of Capacity Gap in Distilling Language Models},
3 author={Zhang, Chen and Song, Dawei and Ye, Zheyu and Gao, Yan},
4 year={2023},
5 url={https://arxiv.org/abs/2311.07052}
6}| Metric | Value |
|---|---|
| Avg. | 42.94 |
| ARC (25-shot) | 44.03 |
| HellaSwag (10-shot) | 67.19 |
| MMLU (5-shot) | 39.17 |
| TruthfulQA (0-shot) | 45.67 |
| Winogrande (5-shot) | 65.27 |
| GSM8K (5-shot) | 10.54 |
| DROP (3-shot) | 28.73 |