import gradio as gr
import requests
import os
 
token = os.getenv("HF_TOKEN")

# Function to generate text
def generate_text(prompt, model_choice, max_tokens, *other_params):
    # Model name matching с URL
    model_urls = {
        "GPT-3.5": "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1",
        "GPT-4": "https://api-inference.huggingface.co/models/ai-forever/ruGPT-4"
    }
    # Selecting a URL depending on the selected model
    api_url = model_urls[model_choice]

    # Preparing data for a request
    headers = {
        "Authorization": f"Bearer {token}"
    }
    payload = {
        "inputs": prompt,
        "parameters": {"max_length": max_tokens},
        "options": {"use_cache": False}
    }

    # Sending a request to API
    response = requests.post(api_url, headers=headers, json=payload)
    if response.status_code == 200:
        # Returning the generated text
        return response.json()[0]['generated_text']
    else:
        # Returning an error message
        return "Error: " + response.text

# Creating an interface using G Blocks
with gr.Blocks() as demo:
    with gr.Tab("Basic settings"):
        with gr.Row():
            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Enter text...")
            model_choice = gr.Radio(["Mixtral-8x7B", "GPT-4"], label="Model selection", value="GPT-3.5")
    
    with gr.Tab("Advanced settings"):
        with gr.Row():
            max_tokens = gr.Slider(100, 5000, step=1, label="Maximum tokens")
            # Here you can add other parameters for the text generation API
    
    with gr.Row():
        generate_btn = gr.Button("Generation")
    
    with gr.Row():
        output_text = gr.Textbox(label="Answer", placeholder="The generated text will be here...")
    
    # Setting up a callback function for the button
    generate_btn.click(
        fn=generate_text,
        inputs=[prompt, model_choice, max_tokens],
        outputs=output_text
    )

# Launching the interface
demo.launch()