import gradio as gr
import os
from huggingface_hub import InferenceClient

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""

model_name = "meta-llama/Llama-3.2-1B"
huggingface_token = os.getenv("SECRET_ENV_VARIABLE")
#client = InferenceClient(api_key=huggingface_token)
client = InferenceClient(model=model_name, token=huggingface_token)


def generate_text(
    prompt,
    system_message,
    max_tokens,
    temperature,
    top_p
):
    try:
        print(f"Attempting to generate text for prompt: {prompt[:50]}...")
        
        response = client.text_generation(
            prompt,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_k=50,
            top_p=top_p,
            do_sample=True
        )
        
        print(f"Generated text: {response[:100]}...")
        return response
    except Exception as e:
        print(f"Error in generate_text: {type(e).__name__}: {str(e)}")
        return f"An error occurred: {type(e).__name__}: {str(e)}"



with gr.Blocks() as demo:
    gr.Markdown("Q&A App")

    with gr.Tab("Q&A"):
        Query = gr.Textbox(label="Query")
        generate_button = gr.Button("Ask Query")
        output = gr.Textbox(label="Generated Answer", lines=10)
        
        generate_button.click(generate_text, 
        #inputs=[industry, recipient_role, company_details], 
            inputs=[
                Query,
                gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-p (nucleus sampling)",
                ),
            ],    
        outputs=output)



if __name__ == "__main__":
    demo.launch()