Update generate.py

inference/generate.py

@@ -12,19 +12,13 @@ from model import Transformer, ModelArgs
 
 
 def sample(logits, temperature: float = 1.0):
-    """
+    try:
-    Samples a token from the logits using temperature scaling.
+        logits = logits / max(temperature, 1e-5)
-
+        probs = torch.softmax(logits, dim=-1)
-    Args:
+        return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1)
-        logits (torch.Tensor): The logits tensor for token predictions.
+    except Exception as e:
-        temperature (float, optional): Temperature for scaling logits. Defaults to 1.0.
+        print(f"Erro ao amostrar token: {e}")
-
+        raise
-    Returns:
-        torch.Tensor: The sampled token.
-    """
-    logits = logits / max(temperature, 1e-5)
-    probs = torch.softmax(logits, dim=-1)
-    return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1)
 
 
 @torch.inference_mode()
@@ -35,47 +29,40 @@ def generate(
     eos_id: int,
     temperature: float = 1.0
 ) -> List[List[int]]:
-    """
+    try:
-    Generates new tokens based on the given prompt tokens using the specified model.
+        prompt_lens = [len(t) for t in prompt_tokens]
+        assert max(prompt_lens) <= model.max_seq_len, f"Prompt é maior do que o comprimento máximo do modelo ({model.max_seq_len})"
+        total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens))
+        tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long, device="cuda")
         
-    Args:
+        for i, t in enumerate(prompt_tokens):
-        model (Transformer): The transformer model used for token generation.
+            tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long, device="cuda")
-        prompt_tokens (List[List[int]]): A list of lists containing the prompt tokens for each sequence.
-        max_new_tokens (int): The maximum number of new tokens to generate.
-        eos_id (int): The end-of-sequence token ID.
-        temperature (float, optional): The temperature value for sampling. Defaults to 1.0.
             
-    Returns:
+        prev_pos = 0
-        List[List[int]]: A list of lists containing the generated tokens for each sequence.
+        finished = torch.tensor([False] * len(prompt_tokens), device="cuda")
-    """
+        prompt_mask = tokens != -1
-    prompt_lens = [len(t) for t in prompt_tokens]
+        
-    assert max(prompt_lens) <= model.max_seq_len
+        for cur_pos in range(min(prompt_lens), total_len):
-    total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens))
+            logits = model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
-    tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long, device="cuda")
+            next_token = sample(logits, temperature) if temperature > 0 else logits.argmax(dim=-1)
-    for i, t in enumerate(prompt_tokens):
+            next_token = torch.where(prompt_mask[:, cur_pos], tokens[:, cur_pos], next_token)
-        tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long, device="cuda")
+            tokens[:, cur_pos] = next_token
-    prev_pos = 0
+            finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id)
-    finished = torch.tensor([False] * len(prompt_tokens), device="cuda")
+            prev_pos = cur_pos
-    prompt_mask = tokens != -1
+            if finished.all():
-    for cur_pos in range(min(prompt_lens), total_len):
+                break
-        logits = model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
+        
-        if temperature > 0:
+        completion_tokens = []
-            next_token = sample(logits, temperature)
+        for i, toks in enumerate(tokens.tolist()):
-        else:
+            toks = toks[prompt_lens[i]:prompt_lens[i] + max_new_tokens]
-            next_token = logits.argmax(dim=-1)
+            if eos_id in toks:
-        next_token = torch.where(prompt_mask[:, cur_pos], tokens[:, cur_pos], next_token)
+                toks = toks[:toks.index(eos_id)]
-        tokens[:, cur_pos] = next_token
+            completion_tokens.append(toks)
-        finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id)
+        
-        prev_pos = cur_pos
+        return completion_tokens
-        if finished.all():
+    except Exception as e:
-            break
+        print(f"Erro ao gerar texto: {e}")
-    completion_tokens = []
+        raise
-    for i, toks in enumerate(tokens.tolist()):
-        toks = toks[prompt_lens[i]:prompt_lens[i]+max_new_tokens]
-        if eos_id in toks:
-            toks = toks[:toks.index(eos_id)]
-        completion_tokens.append(toks)
-    return completion_tokens
 
 
 def main(
@@ -86,100 +73,110 @@ def main(
     max_new_tokens: int = 100,
     temperature: float = 1.0,
 ) -> None:
-    """
+    try:
-    Main function to load the model and perform interactive or batch text generation.
+        if not os.path.exists(ckpt_path):
+            raise FileNotFoundError(f"O caminho para o checkpoint '{ckpt_path}' não existe.")
         
-    Args:
+        if not os.path.exists(config):
-        ckpt_path (str): Path to the model checkpoint directory.
+            raise FileNotFoundError(f"O arquivo de configuração '{config}' não foi encontrado.")
-        config (str): Path to the model configuration file.
-        input_file (str, optional): Path to a file containing input prompts. Defaults to "".
-        interactive (bool, optional): Whether to run in interactive mode. Defaults to True.
-        max_new_tokens (int, optional): Maximum number of new tokens to generate. Defaults to 100.
-        temperature (float, optional): Temperature for sampling. Defaults to 1.0.
-    """
-    world_size = int(os.getenv("WORLD_SIZE", "1"))
-    rank = int(os.getenv("RANK", "0"))
-    local_rank = int(os.getenv("LOCAL_RANK", "0"))
-    if world_size > 1:
-        dist.init_process_group("nccl")
-    global print
-    if rank != 0:
-        print = lambda *_, **__: None
-    torch.cuda.set_device(local_rank)
-    torch.set_default_dtype(torch.bfloat16)
-    torch.set_num_threads(8)
-    torch.manual_seed(965)
-    with open(config) as f:
-        args = ModelArgs(**json.load(f))
-    print(args)
-    with torch.device("cuda"):
-        model = Transformer(args)
-    tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
-    tokenizer.decode(generate(model, [tokenizer.encode("DeepSeek")], 2, -1, 1.)[0])
-    load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
         
-    if interactive:
+        world_size = int(os.getenv("WORLD_SIZE", "1"))
-        messages = []
+        rank = int(os.getenv("RANK", "0"))
-        while True:
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
-            if world_size == 1:
-                prompt = input(">>> ")
-            elif rank == 0:
-                prompt = input(">>> ")
-                objects = [prompt]
-                dist.broadcast_object_list(objects, 0)
-            else:
-                objects = [None]
-                dist.broadcast_object_list(objects, 0)
-                prompt = objects[0]
-            if prompt == "/exit":
-                break
-            elif prompt == "/clear":
-                messages.clear()
-                continue
-            messages.append({"role": "user", "content": prompt})
-            prompt_tokens = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
-            completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id, temperature)
-            completion = tokenizer.decode(completion_tokens[0], skip_special_tokens=True)
-            print(completion)
-            messages.append({"role": "assistant", "content": completion})
-    else:
-        with open(input_file) as f:
-            prompts = [line.strip() for line in f.readlines()]
-        assert len(prompts) <= args.max_batch_size
-        prompt_tokens = [tokenizer.apply_chat_template([{"role": "user", "content": prompt}], add_generation_prompt=True) for prompt in prompts]
-        completion_tokens = generate(model, prompt_tokens, max_new_tokens, tokenizer.eos_token_id, temperature)
-        completions = tokenizer.batch_decode(completion_tokens, skip_special_tokens=True)
-        for prompt, completion in zip(prompts, completions):
-            print("Prompt:", prompt)
-            print("Completion:", completion)
-            print()
         
-    if world_size > 1:
+        if world_size > 1:
-        dist.destroy_process_group()
+            dist.init_process_group("nccl")
+        
+        global print
+        if rank != 0:
+            print = lambda *_, **__: None
+        
+        torch.cuda.set_device(local_rank)
+        torch.set_default_dtype(torch.bfloat16)
+        torch.set_num_threads(8)
+        torch.manual_seed(965)
+        
+        with open(config) as f:
+            args = ModelArgs(**json.load(f))
+        print(args)
+        
+        with torch.device("cuda"):
+            model = Transformer(args)
+        
+        tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
+        
+        tokenizer.decode(generate(model, [tokenizer.encode("DeepSeek")], 2, -1, 1.)[0])
+        
+        load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
+        
+        if interactive:
+            messages = []
+            while True:
+                try:
+                    if world_size == 1:
+                        prompt = input(">>> ")
+                    elif rank == 0:
+                        prompt = input(">>> ")
+                        objects = [prompt]
+                        dist.broadcast_object_list(objects, 0)
+                    else:
+                        objects = [None]
+                        dist.broadcast_object_list(objects, 0)
+                        prompt = objects[0]
+                    
+                    if prompt == "/exit":
+                        break
+                    elif prompt == "/clear":
+                        messages.clear()
+                        continue
+                    
+                    messages.append({"role": "user", "content": prompt})
+                    prompt_tokens = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+                    completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id, temperature)
+                    completion = tokenizer.decode(completion_tokens[0], skip_special_tokens=True)
+                    print(completion)
+                    messages.append({"role": "assistant", "content": completion})
+                
+                except Exception as e:
+                    print(f"Erro durante a interação: {e}")
+                    continue
+        
+        else:
+            with open(input_file) as f:
+                prompts = [line.strip() for line in f.readlines()]
+            
+            assert len(prompts) <= args.max_batch_size
+            prompt_tokens = [tokenizer.apply_chat_template([{"role": "user", "content": prompt}], add_generation_prompt=True) for prompt in prompts]
+            completion_tokens = generate(model, prompt_tokens, max_new_tokens, tokenizer.eos_token_id, temperature)
+            completions = tokenizer.batch_decode(completion_tokens, skip_special_tokens=True)
+            
+            for prompt, completion in zip(prompts, completions):
+                print("Prompt:", prompt)
+                print("Completion:", completion)
+                print()
+    
+    except FileNotFoundError as e:
+        print(f"Erro de arquivo: {e}")
+    
+    except Exception as e:
+        print(f"Erro na execução principal: {e}")
+    
+    finally:
+        if world_size > 1:
+            dist.destroy_process_group()
 
 
 if __name__ == "__main__":
-    """
+    try:
-    Command-line interface for distributed text generation.
+        parser = ArgumentParser()
-
+        parser.add_argument("--ckpt-path", type=str, required=True)
-    Arguments:
+        parser.add_argument("--config", type=str, required=True)
-        --ckpt-path (str): Path to the model checkpoint directory.
+        parser.add_argument("--input-file", type=str, default="")
-        --config (str): Path to the model configuration file.
+        parser.add_argument("--interactive", action="store_true")
-        --input-file (str, optional): File containing prompts for batch processing.
+        parser.add_argument("--max-new-tokens", type=int, default=200)
-        --interactive (bool, optional): Enable interactive mode for generating text.
+        parser.add_argument("--temperature", type=float, default=0.2)
-        --max-new-tokens (int, optional): Maximum number of new tokens to generate. Defaults to 200.
+        args = parser.parse_args()
-        --temperature (float, optional): Temperature for sampling. Defaults to 0.2.
+        assert args.input_file or args.interactive
-
+        main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature)
-    Raises:
+    except Exception as e:
-        AssertionError: If neither input-file nor interactive mode is specified.
+        print(f"Erro na execução do script: {e}")
-    """
-    parser = ArgumentParser()
-    parser.add_argument("--ckpt-path", type=str, required=True)
-    parser.add_argument("--config", type=str, required=True)
-    parser.add_argument("--input-file", type=str, default="")
-    parser.add_argument("--interactive", action="store_true")
-    parser.add_argument("--max-new-tokens", type=int, default=200)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    args = parser.parse_args()
-    assert args.input_file or args.interactive
-    main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature)