| def generate_code(prompt, generator, max_new_tokens=200): | |
| """ | |
| Generate code from a natural language prompt. | |
| """ | |
| output = generator( | |
| prompt, | |
| max_new_tokens=max_new_tokens, | |
| do_sample=True, | |
| top_p=0.9, | |
| temperature=0.7, | |
| eos_token_id=generator.tokenizer.eos_token_id | |
| ) | |
| # Return only new code (without repeating prompt) | |
| return output[0]["generated_text"][len(prompt):].strip() | |