# import torch
# import random

# def sample_word(words, logits, sampling_technique='inverse_transform', temperature=1.0):
#     if sampling_technique == 'inverse_transform':
#         probs = torch.softmax(torch.tensor(logits), dim=-1)
#         cumulative_probs = torch.cumsum(probs, dim=-1)
#         random_prob = random.random()
#         sampled_index = torch.where(cumulative_probs >= random_prob)[0][0]
#     elif sampling_technique == 'exponential_minimum':
#         probs = torch.softmax(torch.tensor(logits), dim=-1)
#         exp_probs = torch.exp(-torch.log(probs))
#         random_probs = torch.rand_like(exp_probs)
#         sampled_index = torch.argmax(random_probs * exp_probs)
#     elif sampling_technique == 'temperature':
#         scaled_logits = torch.tensor(logits) / temperature
#         probs = torch.softmax(scaled_logits, dim=-1)
#         sampled_index = torch.multinomial(probs, 1).item()
#     elif sampling_technique == 'greedy':
#         sampled_index = torch.argmax(torch.tensor(logits)).item()
#     else:
#         raise ValueError("Invalid sampling technique. Choose 'inverse_transform', 'exponential_minimum', 'temperature', or 'greedy'.")
    
#     sampled_word = words[sampled_index]
#     return sampled_word

import torch
import random

def sample_word(sentence, words, logits, sampling_technique='inverse_transform', temperature=1.0):
    if sampling_technique == 'inverse_transform':
        probs = torch.softmax(torch.tensor(logits), dim=-1)
        cumulative_probs = torch.cumsum(probs, dim=-1)
        random_prob = random.random()
        sampled_index = torch.where(cumulative_probs >= random_prob)[0][0]
    elif sampling_technique == 'exponential_minimum':
        probs = torch.softmax(torch.tensor(logits), dim=-1)
        exp_probs = torch.exp(-torch.log(probs))
        random_probs = torch.rand_like(exp_probs)
        sampled_index = torch.argmax(random_probs * exp_probs)
    elif sampling_technique == 'temperature':
        scaled_logits = torch.tensor(logits) / temperature
        probs = torch.softmax(scaled_logits, dim=-1)
        sampled_index = torch.multinomial(probs, 1).item()
    elif sampling_technique == 'greedy':
        sampled_index = torch.argmax(torch.tensor(logits)).item()
    else:
        raise ValueError("Invalid sampling technique. Choose 'inverse_transform', 'exponential_minimum', 'temperature', or 'greedy'.")
    
    sampled_word = words[sampled_index]
    
    # Replace [MASK] with the sampled word
    filled_sentence = sentence.replace('[MASK]', sampled_word)
    
    return filled_sentence