test

research/long-context-llama/H2O/exp.sh

@@ -4,8 +4,13 @@
 # --model-name meta-llama/Llama-2-7b-hf 
 # 20.46/4.9/15.11
 
+# CUDA_VISIBLE_DEVICES=$1 python -u generation.py \
+# --input-path data/summarization/xsum.jsonl \
+# --output-path summarization_output/xsum_h2o.jsonl \
+# --model-name meta-llama/Llama-2-7b-hf \
+# --enable_h2o_generation 
+
 CUDA_VISIBLE_DEVICES=$1 python -u generation.py \
 --input-path data/summarization/xsum.jsonl \
 --output-path summarization_output/xsum_h2o.jsonl \
---model-name meta-llama/Llama-2-7b-hf \
---enable_h2o_generation 
+--model-name meta-llama/Llama-2-7b-hf

research/long-context-llama/H2O/generation.py

@@ -92,7 +92,6 @@ if __name__ == '__main__':
 
             input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors='pt').input_ids.to(model.device)
 
-            print(input_ids)
             output_sequences = model.generate(
                 input_ids=input_ids,
                 max_length=request['max_tokens'] + len(input_ids[0]),
@@ -103,11 +102,6 @@ if __name__ == '__main__':
                 return_dict_in_generate=True, output_scores=True,
             )
 
-            print('Finish')
-
-            # if args.enable_h2o_generation:
-            #     self._clean_cache()
-
             tokens = tokenizer.convert_ids_to_tokens(output_sequences['sequences'].squeeze(0))[len(input_ids[0]):]
             logprobs = [logits.log_softmax(dim=-1).max().item() for logits in output_sequences['scores']]
             top_logprobs = [{i: v for i, v in zip(tokens, logprobs)}]