generator.py

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()