| 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 |
| |
