import gradio as gr
import pandas as pd
from transformers import pipeline
import io

# Load the Hugging Face GPT-Neo model
generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B')

# Function to process CSV file and instruction using Hugging Face GPT-Neo
def process_csv_with_huggingface(file, instruction):
    # Read CSV from file
    df = pd.read_csv(file.name)
    
    # Prepare the prompt for GPT-Neo to generate Python code
    prompt = f"""
You are an expert in Python and data manipulation using Pandas. You will be given a DataFrame and an instruction. 
Interpret the instruction and generate Python code to manipulate the DataFrame accordingly.
Only return the Python code.

Instruction: {instruction}

DataFrame:
{df.head().to_dict()}
"""
    
    # Request code generation from GPT-Neo
    response = generator(prompt, max_length=200, num_return_sequences=1)

    # Extract the generated Python code
    generated_code = response[0]['generated_text'].strip()

    # Execute the generated code
    try:
        exec(generated_code, globals(), locals())
        # After execution, the manipulated DataFrame is stored in 'df'
        # Convert the DataFrame to CSV for download
        output = io.StringIO()
        df.to_csv(output, index=False)
        output.seek(0)
        return output.getvalue()
    except Exception as e:
        return f"Error while processing: {str(e)}"

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# LLM-Powered CSV Manipulation Tool")
    file_input = gr.File(label="Upload CSV File", type="filepath")  # Changed to 'filepath'
    instruction_input = gr.Textbox(label="Instruction (e.g., 'Add a new column Age with value 30')")
    output_file = gr.File(label="Download Manipulated CSV")
    run_button = gr.Button("Apply Instruction")

    # Link inputs and outputs
    run_button.click(process_csv_with_huggingface, inputs=[file_input, instruction_input], outputs=[output_file])

# Launch Gradio Interface
demo.launch()