Update utils.py

utils.py

@@ -53,6 +53,33 @@ def generate_prompt_with_history(text, history, tokenizer, max_length=2048):
     else:
         return None
 
+# Few-Shot Training wird durch ein spezielles Modell-Laden gemacht (setFit)
+def load_tokenizer_and_model_setFit(base_model, load_8bit=False):
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    
+    
+    tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast = True, use_auth_token=True, bos_token='<|startoftext|>', eos_token='<|endoftext|>', pad_token='<|pad|>')
+    if device == "cuda":
+        model = SetFitModel.from_pretrained(
+            base_model,
+            load_in_8bit=load_8bit,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            use_auth_token=True,
+        )
+    else:
+        model = SetFitModel.from_pretrained(
+            base_model, device_map={"": device}, low_cpu_mem_usage=True
+        )
+
+    #if not load_8bit:
+        #model.half()  # seems to fix bugs for some users.
+
+    model.eval()
+    return tokenizer,model, device
 
 
 def load_tokenizer_and_model(base_model, load_8bit=False):