chai _layer for all tasks

[rippleD030]

Could you please clarify whether the number of clusters you chose for each layer is the same across all five tasks?

[iidsample]

Yes we keep the number of clusters across all five tasks same.

[rippleD030]

Hello,
I am testing the performance of the LLaMA 7B model on the PIQA task using the lm-evaluation-harness with the following settings:
ModelArgs(dim=4096, n_layers=32, n_heads=32, vocab_size=32000, multiple_of=256, norm_eps=1e-06, prune_layer=0, max_batch_size=32, max_seq_len=512, chai_activate=True, che=True, chai_layers=[28, 28, 28, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 8, 8, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4])
However, I am getting an accuracy of only 52, which is significantly lower than the 77 mentioned in the paper. Could you please advise if I might have missed any steps in this evaluation process?

[iidsample]

Interesting. I am happy to take a look if you share the code. We have not evaluated this code with LM harness. But ideally the results should be the same.

[rippleD030]

Interesting. I am happy to take a look if you share the code. We have not evaluated this code with LM harness. But ideally the results should be the same.

@iidsample Thank you for your reply. For testing local models with lm-evaluation-harness, it's necessary to implement the loglikelihood method. Below is my implementation on the LLaMA model, but currently, it only accepts test data with a batch size of 1. Could you please let me know if this will impact Chai's performance?

    def __init__(self, llama_model: LLaMA):
        super().__init__()  # Initialize the parent class
        self.model = llama_model  # LLaMA model instance
        self.tokenizer = llama_model.tokenizer  # Get the tokenizer

    def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]:  # Modified to batch processing
        results = []
        for request in requests:
            input_text, target_text = request.args  # Get input text and target text
            input_tokens = self.tokenizer.encode(input_text, bos=True, eos=False)  # Encode the input text
           
            input_tensor = torch.tensor(input_tokens).unsqueeze(0).cuda()  # Add batch dimension
            target_tokens = self.tokenizer.encode(target_text, bos=False, eos=False)  # Encode the target text
            target_tensor = torch.tensor(target_tokens).unsqueeze(0).cuda()  # Add batch dimension
            combined_tokens = torch.cat((input_tensor, target_tensor), dim=-1)  # Concatenate along the last dimension

            # Calculate the length of the input text
            input_len = len(input_tokens)

            log_likelihood = 0.0
            is_greedy = True

            # Calculate the conditional probability for each target token
            prev_pos = 0
            for token in target_tokens:
                # Calculate logits at the current step
                logits = self.model.model.forward(combined_tokens[:, prev_pos:input_len], prev_pos)

                # Get the probability distribution of the current token
                probs = torch.softmax(logits[:, :], dim=-1)

                # Get the log probability of the current token
                token_log_prob = torch.log(probs[0, token])

                log_likelihood += token_log_prob.item()  # Accumulate log likelihood

                # Check if the current token is the greedy result
                if torch.argmax(probs[0, :]) != token:
                    is_greedy = False

                # Update the input sequence to include the current generated token
                prev_pos = input_len
                input_len += 1  # Update input length

            results.append((log_likelihood, is_greedy))

        return results

[rippleD030]

I tried modifying the code to run according to batch size, but the testing results on Chai are still not ideal (I commented out line 196 in llama/model.py where assert num_examples == 1 during the Chai test). Could you please provide any methods or suggestions for testing on Chai? greatly appreciate your response.

def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]:
    results = []
    batch_size = 32

    # Split requests according to batch_size
    for i in range(0, len(requests), batch_size):
        batch_requests = requests[i:min(i + batch_size, len(requests))]

        input_tensors = []
        target_tensors = []
        combined_tensors = []
        input_lengths = []

        for request in batch_requests:
            input_text, target_text = request.args
            input_tokens = self.tokenizer.encode(input_text, bos=True, eos=False)
            if len(input_tokens) > 500:
                input_tokens = self.tokenizer.encode(request.doc['passage'], bos=True, eos=False)[:450] + \
                               self.tokenizer.encode(request.doc['question'], bos=False, eos=False)
            input_tensor = torch.tensor(input_tokens).cuda()
            target_tokens = self.tokenizer.encode(target_text, bos=False, eos=False)
            target_tensor = torch.tensor(target_tokens).cuda()
            combined_tensor = torch.cat((input_tensor, target_tensor), dim=-1)

            input_tensors.append(input_tensor)
            target_tensors.append(target_tensor)

            combined_tensors.append(combined_tensor)

        bsz = len(combined_tensors)
        lengths = [tensor.size(0) for tensor in combined_tensors]  # Length of each question
        input_lengths = [tensor.size(0) for tensor in input_tensors]

        prompt_tokens = combined_tensors
        min_prompt_size = min([len(t) for t in input_tensors])
        max_prompt_size = max([len(t) for t in combined_tensors])
        assert min_prompt_size >= 6
        total_len = max_prompt_size

        tokens = torch.full((bsz, total_len), 0).cuda().long()
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = t.clone().detach().long()

        log_likelihood = torch.zeros(bsz).to("cuda")
        start_pos = min_prompt_size
        prev_pos = 0
        for cur_pos in range(start_pos, total_len):
            logits = self.model.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
            probs = torch.softmax(logits, dim=-1)
            loc = tokens[:, cur_pos]

            # Assume probs is a list where each element is a tensor
            token_log_prob = torch.log(torch.tensor([prob[loc[i]] for i, prob in enumerate(probs)])).to("cuda")

            # Only keep the output position's values for loglikelihood
            for i, length in enumerate(lengths):
                if length < cur_pos:
                    token_log_prob[i] = 0

            for i, length in enumerate(input_lengths):
                if length > cur_pos:
                    token_log_prob[i] = 0

            log_likelihood += token_log_prob

            is_greedy = False
            prev_pos = cur_pos
        results.extend([(l.item(), False) for l in log_likelihood])

    return results

[iidsample]

The batch size should not have any affect on accuracy. I will try to setup lm evaluation harness locally and see what is happening. Thanks.