import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load Hugging Face model and tokenizer
model_name = "abrotech/Zora-ALM-7.2B-gguf"  # Your Hugging Face model space
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define function to handle user input and generate response
def generate_response(user_input):
    inputs = tokenizer(user_input, return_tensors="pt")
    outputs = model.generate(input_ids=inputs["input_ids"], max_length=150, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Set up the Gradio interface
with gr.Blocks() as demo:
    gr.HTML("<h1 style='text-align: center;'>Welcome to Zora Assistant</h1>")
    gr.HTML("<p style='text-align: center;'>Ask anything and Zora will answer!</p>")

    with gr.Row():
        with gr.Column():
            user_input = gr.Textbox(label="Enter your question", placeholder="Ask Zora anything...")
            submit_btn = gr.Button("Get Answer")
            response_output = gr.Textbox(label="Zora's Answer", interactive=False)

    submit_btn.click(generate_response, inputs=user_input, outputs=response_output)

# Launch the Gradio app
demo.launch(share=True)