import gradio as gr
from huggingface_hub import InferenceClient

# Use a model that supports text-generation
client = InferenceClient("meta-llama/Llama-3.2-1B-Instruct")

def generate(prompt, temperature=0.8, max_tokens=256):
    try:
        response = client.text_generation(
            prompt=prompt,
            temperature=temperature,
            max_new_tokens=max_tokens,
            do_sample=True if temperature > 0 else False
        )
        return response
    except Exception as e:
        return f"Error: {str(e)}"

with gr.Blocks(title="amkyaw-coder") as demo:
    gr.Markdown("# amkyaw-coder\n🤖 Code Generation Model (via HF Inference)")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=4, placeholder="Enter your prompt here...")
            temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature")
            max_tokens = gr.Slider(32, 512, value=128, step=32, label="Max Tokens")
            submit = gr.Button("Generate", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(label="Output", lines=15)
    
    submit.click(generate, inputs=[prompt, temperature, max_tokens], outputs=output)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)