class Node:
def __init__(self, freq):
self.left = None
self.right = None
self.father = None
self.freq = freq
def isLeft(self):
return self.father.left == self
def createNodes(freqs):
return [Node(freq) for freq in freqs]
def createHuffmanTree(nodes):
queue = nodes[:]
while len(queue) > 1:
queue.sort(key=lambda item:item.freq)
node_left = queue.pop(0)
node_right = queue.pop(0)
node_father = Node(node_left.freq + node_right.freq)
node_father.left = node_left
node_father.right = node_right
node_left.father = node_father
node_right.father = node_father
queue.append(node_father)
queue[0].father = None
return queue[0]
def huffmanEncoding(nodes, root):
codes = [''] * len(nodes)
for i in range(len(nodes)):
node_tmp = nodes[i]
while node_tmp != root:
if node_tmp.isLeft():
codes[i] = '0' + codes[i]
else:
codes[i] = '1' + codes[i]
node_tmp = node_tmp.father
return codes
class Huffman():
def __init__(self, k=None, bits=None):
self.k = k
self.bits = bits
def extract(self, pred_prob, word_index):
bit = None
prob, idx = torch.topk(pred_prob, k=self.k, dim=1)
word_prob = [[i.item(),j.item()] for i,j in zip(idx[0], prob[0])]
nodes = createNodes([item[1] for item in word_prob])
root = createHuffmanTree(nodes)
codes = huffmanEncoding(nodes, root)
for i, w_p in enumerate(word_prob):
if w_p[0] == word_index:
bit = codes[i]
break
return bit
def __call__(self, pred_prob):
prob, idx = torch.topk(pred_prob, k=self.k, dim=1)
word_prob = [[i.item(),j.item()] for i,j in zip(idx[0], prob[0])]
nodes = createNodes([item[1] for item in word_prob])
root = createHuffmanTree(nodes)
codes = huffmanEncoding(nodes, root)
for i,code in enumerate(codes):
bit = self.bits[:len(code)]
if bit == code:
bit_word_index = word_prob[i][0]
self.bits = self.bits[len(code):]
break
if self.extract(pred_prob,bit_word_index)!=bit:
print("False")
bit_word_index = torch.LongTensor([bit_word_index]).to(pred_prob.device)
return bit_word_index
class FLC():
def __init__(self, k=None, bits=None):
self.k = k
self.bits = bits
def extract(self, pred_prob, word_index):
bit = None
prob, idx = torch.topk(pred_prob, k=2**self.k, dim=1)
word_prob = [[i.item(),j.item()] for i,j in zip(idx[0], prob[0])]
codes = [str(bin(i))[2:].zfill(self.k) for i in range(2**self.k)]
for i, w_p in enumerate(word_prob):
if w_p[0] == word_index:
bit = codes[i]
break
return bit
def __call__(self, pred_prob):
prob, idx = torch.topk(pred_prob, k=2**self.k, dim=1)
word_prob = [[i.item(),j.item()] for i,j in zip(idx[0], prob[0])]
codes = [str(bin(i))[2:].zfill(self.k) for i in range(2**self.k)]
bit = self.bits[:self.k]
for i,code in enumerate(codes):
if bit==code:
bit_word_index = word_prob[i][0]
self.bits = self.bits[self.k:]
break
bit_word_index = torch.LongTensor([bit_word_index]).to(pred_prob.device)
return bit_word_index
def setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
def generate_text_with_steganography(model, tokenizer, prompt, handle, max_length=30):
device = model.device
input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
# Generate steganographic text
stega_output = input_ids.clone()
normal_output = input_ids.clone()
for _ in range(max_length):
if len(handle.bits) == 0:
break
# Get model predictions
outputs = model(stega_output)
next_token_logits = outputs.logits[:, -1, :]
next_token_probs = F.softmax(next_token_logits, dim=-1)
# Apply steganography
bit_word_index = handle(next_token_probs)
stega_output = torch.cat([stega_output, bit_word_index.unsqueeze(1)], dim=1)
# Generate normal text
outputs = model(normal_output)
next_token_logits = outputs.logits[:, -1, :]
next_token_probs = F.softmax(next_token_logits, dim=-1)
next_token = torch.multinomial(next_token_probs, num_samples=1)
normal_output = torch.cat([normal_output, next_token], dim=1)
if next_token[0].item() == tokenizer.eos_token_id:
break
return stega_output, normal_output