Update app.py

app.py

@@ -57,14 +57,12 @@ sample_input = tok.apply_chat_template(chat, tokenize=False, add_generation_prom
 input_start_id=sample_input.find(slot)
 prefix=sample_input[:input_start_id]
 suffix=sample_input[input_start_id+len(slot):]
-print(tok.encode(prefix,return_tensors="pt")[0])
-print(tok.encode(suffix,return_tensors="pt")[0])
 prefix_embedding=embedding_func(
     tok.encode(prefix,return_tensors="pt")[0]
 )
 suffix_embedding=embedding_func(
     tok.encode(suffix,return_tensors="pt")[0]
-)
+)[1:]
 
 #print(prefix_embedding)
 print(f"Sucessfully loaded the model to the memory")
@@ -84,7 +82,7 @@ def embedding_shift(original_embedding,shift_embeddings,prefix_embedding,suffix_
     return input_embeddings
 def engine(input_embeds):
     output_text = []
-    batch_size = 10
+    batch_size = 5
     with torch.no_grad():
         for start in range(0,len(input_embeds),batch_size):
             batch_input_embeds = input_embeds[start:start+batch_size]
@@ -175,7 +173,7 @@ def chat(message, history, with_defense,perturb_times):
     generate_kwargs = dict(
         model_inputs,
         streamer=streamer,
-        max_new_tokens=1024,
+        max_new_tokens=256,
         do_sample=True,
         top_p=0.90,
         temperature=0.6,