Function process

evaluation/generation.py:195–243

(obj)

    )
    
    def process(obj):
        results = []
        prompt = obj["prompt"]
        if args.generation_mode == "instruction":
            inputs = tokenizer([prompt] * args.micro_batch_size, return_tensors="pt")
            inputs = inputs.to(model.device)
            outputs = model.generate(**inputs,
                                    max_length=args.max_length,
                                    do_sample=True if not args.greedy else False,
                                    use_cache=True,
                                    top_p=args.top_p,
                                    top_k=args.top_k,
                                    temperature=args.temperature,
                                    pad_token_id=tokenizer.eos_token_id)
            for i, output in enumerate(outputs):
                response = tokenizer.decode(output)
                res = obj.copy()
                res["generation"] = response[len(prompt):].strip()
                results.append(res)
        elif args.generation_mode == "completion":
            inputs = tokenizer([prompt for _ in range(args.micro_batch_size)], return_tensors="pt")
            inputs = inputs.to(model.device)
            stop_criteria = CodeStoppingCriteria(
                max_length=args.max_length,
                micro_batch_size=args.micro_batch_size,
                tokenizer=tokenizer,
                dataset_type=args.dataset_type,
                language_type=args.language_type,
                prompt=prompt)
            outputs = model.generate(**inputs,
                                    max_length=args.max_length,
                                    do_sample=True if not args.greedy else False,
                                    use_cache=True,
                                    stopping_criteria=[stop_criteria],
                                    top_p=args.top_p,
                                    top_k=args.top_k,
                                    temperature=args.temperature,
                                    pad_token_id=tokenizer.eos_token_id)
            for i, output in enumerate(outputs):
                response = tokenizer.decode(output)
                res = obj.copy()
                res["generation_raw"] = response
                res["generation"] = cleanup_code(
                    response[len(prompt):], 
                    dataset_type=args.dataset_type,
                    language_type=args.language_type)
                results.append(res)
        
        return results
    
    fout = open(output_path, "w", encoding="utf-8")
    while True: