import gradio as gr
from transformers import pipeline
import torch
import subprocess
import spaces
import os



# Initialize the model pipeline
generator = pipeline('text-generation', model='Locutusque/ChatHercules-2.5-Mistral-7B-DPO', torch_dtype=torch.bfloat16, token=os.environ["HF"])
@spaces.GPU
def generate_text(prompt, temperature, top_p, top_k, repetition_penalty, max_length):
    # Generate text using the model
    generator.model.cuda()
    generator.device = torch.device("cuda")
    prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
    outputs = generator(
        prompt, 
        do_sample=True,
        max_new_tokens=max_length,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        repetition_penalty=repetition_penalty,
        return_full_text=False
    )
    # Extract the generated text and return it
    generated_text = outputs[0]['generated_text']
    return generated_text
# Create the Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Prompt", lines=2, placeholder="Type a prompt..."),
        gr.Slider(minimum=0.1, maximum=2.0, step=0.01, value=0.8, label="Temperature"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.95, label="Top p"),
        gr.Slider(minimum=0, maximum=100, step=1, value=40, label="Top k"),
        gr.Slider(minimum=1.0, maximum=2.0, step=0.01, value=1.10, label="Repetition Penalty"),
        gr.Slider(minimum=5, maximum=4096, step=5, value=1024, label="Max Length")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="ChatHercules-2.5-Mistral-7B",
    description="Try out the ChatHercules-2.5-Mistral-7B model for free!"
)

iface.launch()