Spaces:
Runtime error
Runtime error
File size: 1,653 Bytes
19cef65 56006e4 19cef65 56006e4 9547676 b45d53a 56006e4 19cef65 9a6b8a0 b45d53a 9a6b8a0 19cef65 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import os
import random
from threading import Thread
from typing import Iterable
import torch
from huggingface_hub import HfApi
from datasets import load_dataset
from transformers import T5Tokenizer, T5ForConditionalGeneration
from transformers import AutoTokenizer, AutoModelForCausalLM
TOKEN = os.environ.get("HF_TOKEN", None)
type2dataset = {
"re2text-easy": load_dataset('3B-Group/ConvRe', "en-re2text", token=TOKEN, split="prompt1"),
"re2text-hard": load_dataset('3B-Group/ConvRe', "en-re2text", token=TOKEN, split="prompt4"),
"text2re-easy": load_dataset('3B-Group/ConvRe', "en-text2re", token=TOKEN, split="prompt1"),
"text2re-hard": load_dataset('3B-Group/ConvRe', "en-text2re", token=TOKEN, split="prompt3")
}
model_id = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id, token=TOKEN)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, token=TOKEN)
# type2dataset = {}
def generate(input_text, sys_prompt) -> str:
sys_prompt = f'''[INST] <<SYS>>
{sys_prompt}
<</SYS>>
'''
input_str = sys_prompt + input_text + " [/INST]"
input_ids = tokenizer(input_str, return_tensors="pt", torch_dtype=torch.bfloat16).input_ids
outputs = model.generate(input_ids, max_length=512)
result = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
return result
def random_examples(dataset_key) -> str:
# target_dataset = type2dataset[f"{task.lower()}-{type.lower()}"]
target_dataset = type2dataset[dataset_key]
idx = random.randint(0, len(target_dataset) - 1)
item = target_dataset[idx]
return item['query']
|