File size: 2,982 Bytes
484dc56 |
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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import torch
import numpy as np
base_model_id = 'teknium/OpenHermes-2.5-Mistral-7B'
ft_model_id = 'mamakos/CMClassifier'
class classifier:
def __init__(self):
bnb_config = BitsAndBytesConfig(load_in_4bit=True,
bnb_4bit_quant_type='nf4',
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=False)
base_model = AutoModelForCausalLM.from_pretrained(base_model_id,
quantization_config=bnb_config,
device_map={'': 'cuda'},
trust_remote_code=True)
self.tokenizer = AutoTokenizer.from_pretrained(base_model_id,
add_bos_token=True,
trust_remote_code=True)
self.ft_model = PeftModel.from_pretrained(base_model, ft_model_id)
self.ft_model.eval()
def predict(self, texts):
if isinstance(texts, str):
texts = [texts]
prompts = [self.get_prompt(text) for text in texts]
probs = np.zeros(len(texts))
with torch.no_grad():
for i in range(len(texts)):
tokenized_prompt = self.tokenizer(prompts[i], return_tensors='pt').to('cuda')
output = self.ft_model.generate(**tokenized_prompt,
pad_token_id=self.tokenizer.eos_token_id,
max_new_tokens=1,
do_sample=False,
return_dict_in_generate=True,
output_scores=True)
transition_score = self.ft_model.compute_transition_scores(output.sequences,
output.scores,
normalize_logits=True)
prob = np.exp(transition_score[0][0].cpu().numpy())
response = self.tokenizer.decode(output.sequences[:, -1][0])
if response == 'No':
prob = 1 - prob
probs[i] = prob
return probs
def get_prompt(self, text):
prompt = f'''<|im_start|>system
You are a helpful assistant.
<|im_end|>
<|im_start|>user
Classify this text as to whether it displays closed-mindedness: "{text}"
If this text displays closed-mindedness, your response must be "Yes".
If this text doesn\'t display closed-mindedness, your response must be "No".
<|im_end|>
<|im_start|>assistant
'''
return prompt
|