evilfreelancer commited on
Commit
2b4daee
·
verified ·
1 Parent(s): 18e2c8e

Create test.py

Browse files
Files changed (1) hide show
  1. test.py +23 -0
test.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
3
+
4
+ MODEL_NAME = "evilfreelancer/llama2-7b-toxicator-ru"
5
+ DEFAULT_INSTRUCTION = "Перефразируй нетоксичный текст так, чтобы он стал токсичным, сохраняя при этом исходный смысл, орфографию и пунктуацию."
6
+ DEFAULT_TEMPLATE = "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
7
+
8
+ # Init model and tokenizer
9
+ generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
10
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
11
+ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
12
+ model.eval()
13
+
14
+ # Build instruct prompt
15
+ user_message = "Великолепный полёт мысли, сразу видно, что Вы очень талантливы."
16
+ prompt = DEFAULT_TEMPLATE.format(**{"instruction": DEFAULT_INSTRUCTION, "input": user_message})
17
+
18
+ # Run model
19
+ data = tokenizer(prompt, return_tensors="pt")
20
+ data = {k: v.to(model.device) for k, v in data.items()}
21
+ output_ids = model.generate(**data, max_length=256, generation_config=generation_config)[0]
22
+ output = tokenizer.decode(output_ids, skip_special_tokens=True)
23
+ print(output)