import os
from huggingface_hub import InferenceClient

client = InferenceClient(
    model="google/flan-t5-large",
    token=os.getenv("HUGGINGFACEHUB_API_TOKEN")
)

def generate_code(prompt):
    full_prompt = f"""You are a professional Python Pandas developer.
Write a complete, executable, valid Python Pandas script to perform the following task on a DataFrame named df:
Task: {prompt}
Make sure the code can be executed without any syntax errors.
Output only valid Python code. No explanation, no comments."""

    response = client.text_generation(
        prompt=full_prompt,
        max_new_tokens=200,     # Updated token limit
        temperature=0.2,
        top_p=0.9,
        repetition_penalty=1.05
    )
    return response