import gradio as gr
from transformers import pipeline

model = pipeline("text-generation", model="umarbutler/open-australian-legal-llm")
    
def chat_fn(message, history):
    prompt = """
        You are a helpful, friendly and very senior lawyer who is an expert in criminal law advocacy and procedure.
        You advise clients on how to deal with police, police powers and what to do if arrested. 
        You also represent in court and know everything about evidence and procedure. Please provide advice to your client.
        Do not guess. Research the relevant legislation, then case law if relevant. Ask quesiton if you need clarification
        for example, which jurisdiction, or about the circumstances. Here is your client's message: 
        
        """ + message
    
    # Generate text using the model
    response = model(prompt, max_length=500, num_return_sequences=1)
    
    # Extract the generated text from the response
    reply = response[0]['generated_text']
    
    # Update history with the new message format
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": reply})
    return history

with gr.Blocks() as dennis_denuto:
    chatbot = gr.Chatbot(
        placeholder="<strong>It's the vibe.</strong><br>Ask anything about criminal law and procedure in Australia.",
        type="messages"
    )
    gr.ChatInterface(
        fn=chat_fn, 
        type="messages", 
        chatbot=chatbot
    )
    
dennis_denuto.launch()