import gradio as gr
from transformers import pipeline

raven_pipeline = pipeline(
    "text-generation",
    model="Nexusflow/NexusRaven-V2-13B",
    torch_dtype="auto",
    device_map="auto",
)

class DialogueToSpeechConverter:
    def __init__(self):
        self.raven_pipeline = raven_pipeline

    def process_text(self, input_text: str) -> str:
        prompt = f"User Query: {input_text}<human_end>"
        result = self.raven_pipeline(prompt, max_new_tokens=2048, return_full_text=False, do_sample=False, temperature=0.001)[0]["generated_text"]
        return result

# Gradio interface
def create_interface():
    converter = DialogueToSpeechConverter()
    with gr.Blocks() as app:
        gr.Markdown("""# 🙋🏻‍♂️Welcome to🌟Tonic's Nexus🐦‍⬛Raven""")
        gr.Markdown("""You can build with this endpoint using Nexus Raven. The demo is still a work in progress but we hope to add some endpoints for commonly used functions such as intention mappers and audiobook processing.""")
        with gr.Row():
            input_text = gr.Textbox(label="Input Text")
            output_text = gr.Textbox(label="Nexus🐦‍⬛Raven", readonly=True)
        input_text.change(converter.process_text, inputs=input_text, outputs=output_text)

    return app

if __name__ == "__main__":
    app = create_interface()
    app.launch()