Upload closedmindeness.py

closedmindeness.py

@@ -0,0 +1,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
+