import gradio as gr
import openai

# Set up OpenAI API key
openai.api_key = "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"

# Define the code generation function
def code_generation(message):
    # Set the model name and parameters
    model = "nlpxucan/wizardlm-base"
    max_tokens = 100

    # Generate code using the model
    response = openai.Completion.create(
        engine="text-davinci-003",
        prompt=message,
        max_tokens=max_tokens,
        temperature=0.7,
        top_p=1.0,
        n=1,
        stop=None,
        frequency_penalty=0.0,
        presence_penalty=0.0,
        log_level="info",
        model=model
    )

    return response.choices[0].text.strip()

# Create the Gradio interface
iface = gr.Interface(
    fn=code_generation,
    inputs="text",
    outputs="text",
    title="Code Generation with WizardLM",
    description="Enter a message as code prompt to generate code.",
    examples=[
        ["Create a function to calculate the factorial of a number."],
        ["Sort a list of integers in ascending order."],
    ],
    theme="default"
)

# Run the Gradio interface
iface.launch()