Generate updates (#8)

* updated generate.py

---------

Co-authored-by: Zhaojun Xie <zxie12@login-1.zaratan.umd.edu>
Co-authored-by: Daniel Nichols <dando18studios@gmail.com>

Author: 3 people

generate.py

@@ -10,7 +10,7 @@
 parser.add_argument('--output', help='Path to the output JSON file')
 parser.add_argument('--max_new_tokens', type=int, default=500, help='Maximum number of new tokens to generate (default: 500)')
 parser.add_argument('--max_length', type=int, default=100, help='Maximum length (default: 500)')
-parser.add_argument('--num-samples-per-prompt', type=int, default=100, help='Number of code samples to generate (default: 100)')
+parser.add_argument('--num_samples_per_prompt', type=int, default=100, help='Number of code samples to generate (default: 100)')
 parser.add_argument('--temperature', type=float, default=0.9, help='Temperature for controlling randomness (default: 0.9)')
 parser.add_argument('--top_p', type=float, default=0.9, help='Top p value for nucleus sampling (default: 0.9)')
 parser.add_argument('--do_sample', action='store_true', help='Enable sampling (default: False)')
@@ -21,17 +21,33 @@
 
 generator = pipeline(model=args.model, torch_dtype=torch.float16, device_map="auto")
 
-responses = {}
+responses = []
 
 for item in prompts:
+    answer = {}
+    problem_type = item["problem_type"]
+    language = item["language"]
     name = item["name"]
+    paral_model = item["parallelism_model"]
     prompt = item["prompt"]
 
-    print(f"Generating code for: {name}")
-    
-    response = generator(prompt, max_length=args.max_length, max_new_tokens=args.max_new_tokens, do_sample=args.do_sample, temperature=args.temperature, top_p=args.top_p)
+    answer["problem_type"] = problem_type
+    answer["language"] = language
+    answer["name"] = name
+    answer["parallelism_model"] = paral_model
+    answer["prompt"] = prompt
+    output = []
 
-    responses[name] = response[0]['generated_text']
+    for x in range(args.num_samples_per_prompt):
+        print(f"Generating code for: {name}")
+
+        response = generator(prompt, max_length=args.max_length, max_new_tokens=args.max_new_tokens, do_sample=args.do_sample, temperature=args.temperature, top_p=args.top_p)
+
+        output.append(response[0]['generated_text'].replace(prompt, "").strip())
+
+    answer["output"] = output
+
+    responses.append(answer)
 
 if args.output:
     with open(args.output, 'w') as output_file: