app.py

import gradio as gr
import os
import torch

from transformers import GemmaTokenizer, AutoModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread


DESCRIPTION = '''
<div>
<h1 style="text-align: center;">Indonesian Legal Assistant</h1>
</div>
'''

LICENSE = """
<p/>
---
Built with Deepseek Qwen3
"""

PLACEHOLDER = """
<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">DeepSeek-R1-0528-Qwen3-8B</h1>
   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Finetuned with Indonesian Legal Consulting Data</p>
</div>
"""


css = """
h1 {
  text-align: center;
  display: block;
}
#duplicate-button {
  margin: auto;
  color: white;
  background: #1565c0;
  border-radius: 100vh;
}
"""

# Load the tokenizer and model
model_name = 'Azzindani/Deepseek_ID_Legal_Preview'

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code = True)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map = 'auto', trust_remote_code = True, torch_dtype = torch.float16)
#model.to('cuda' if torch.cuda.is_available() else 'cpu')

def chat(message, history, temperature = 0.7, max_new_tokens = 1024, show_thinking = True):
    SYSTEM_PROMPT = """
    Anda adalah asisten AI yang ahli di bidang hukum Indonesia. Anda dapat membantu konsultasi hukum, menjawab pertanyaan, dan memberikan analisis berdasarkan peraturan perundang-undangan yang relevan.
    Untuk setiap respons, Anda harus berfikir dan menjawab dengan Bahasa Indonesia, serta gunakan format:
    <think>
    ...
    </think>
    Tuliskan jawaban akhir secara jelas, ringkas, profesional, dan berempati jika diperlukan. Gunakan bahasa hukum yang mudah dipahami. Sertakan referensi hukum Indonesia yang relevan. Selalu rekomendasikan konsultasi dengan ahli hukum untuk keputusan final.
    """
    
    # Prepare conversation history
    conversation = [{'role' : 'system', 'content' : SYSTEM_PROMPT}]
    for user_msg, assistant_msg in history:
        conversation.append({'role' : 'user', 'content' : user_msg})
        conversation.append({'role' : 'assistant', 'content' : assistant_msg})
    conversation.append({'role' : 'user', 'content' : message})
    
    # Ensure tokenizer and model are properly referenced
    try:
        input_ids = tokenizer.apply_chat_template(
            conversation,
            tokenize = True,
            add_generation_prompt = True,
            return_tensors = 'pt'
        ).to(model.device)
    except Exception as e:
        return f"Error preparing input: {str(e)}"
        
    streamer = TextIteratorStreamer(tokenizer, skip_prompt = True, skip_special_tokens = True)
    generate_kwargs = {
        'input_ids' : input_ids,
        'streamer' : streamer,
        'max_new_tokens' : max_new_tokens,
        'do_sample' : temperature > 0,
        'temperature' : temperature if temperature > 0 else 1.0,
        'min_p' : 0.1,
        'top_p' : 1.0,
        'top_k' : 20,
        #'repetition_penalty' : 1.1
    }
    
    # Start generation in a thread to avoid blocking
    thread = Thread(target = model.generate, kwargs = generate_kwargs)
    thread.start()
    
    # Buffers and flags
    thinking_content = []
    final_answer = []
    live_output = []
    in_thinking_block = False
    saw_think_tag = False
    thinking_header_shown = False
    
    # Tracking streaming state
    has_started_output = False
    accumulated_text = ''
    
    for new_text in streamer:
        accumulated_text += new_text
        has_started_output = True
        
        # Detect thinking tags
        if '<think>' in new_text:
            in_thinking_block = True
            saw_think_tag = True
            new_text = new_text.replace('<think>', '')
            if not thinking_header_shown:
                live_output.append('\n🧠 **Sedang berfikir...**\n')
                thinking_header_shown = True
        
        if '</think>' in new_text:
            in_thinking_block = False
            new_text = new_text.replace('</think>', '')
            live_output.append('\n\n-----\n✅ **Sedang menjawab...**\n')
        
        # Handle output based on current state
        if saw_think_tag:
            # Normal operation with thinking tags
            if in_thinking_block:
                thinking_content.append(new_text)
                live_output.append(new_text)
            else:
                # After thinking block, everything is the answer
                final_answer.append(new_text)
                live_output.append(new_text)
        else:
            # No thinking tags detected yet
            if len(accumulated_text) > 20 and not saw_think_tag:
                # No thinking format detected, treat everything as answer
                if not thinking_header_shown:
                    live_output.append("\n⏭️ **Jawaban langsung:**\n\n")
                    thinking_header_shown = True
                
                final_answer.append(new_text)
                live_output.append(new_text)
            else:
                # Still accumulating to detect format
                live_output.append(new_text)
        
        # Always yield current output
        yield ''.join(live_output)
    
    # Final output format handling
    if thinking_content:
        # Show collapsible thinking section if there was thinking content
        final_output = (
            '<details><summary>🧠 <b>Proses berfikir (klik untuk melihat)</b></summary>\n\n'
            + ''.join(thinking_content) +
            '\n</details>\n\n'
            + '-----\n✅ **Jawaban:**\n'
            + ''.join(final_answer)
        )
        yield final_output
    
    # If no final answer was produced but we have output, use the entire output
    if not final_answer and has_started_output:
        final_answer = [accumulated_text]
    
    # Update chat history with the clean answer only (no thinking/formatting)
    history.append((message, ''.join(final_answer)))
    return None  # Explicit return to end function after yielding
        
# Gradio block
chatbot = gr.Chatbot(height = 500, placeholder = PLACEHOLDER, label = 'Legal Assistant')

with gr.Blocks(fill_height = True, css = css) as demo:
    
    gr.Markdown(DESCRIPTION)
    #gr.DuplicateButton(value = "Duplicate Space for private use", elem_id = "duplicate-button")
    gr.ChatInterface(
        fn = chat,
        chatbot = chatbot,
        fill_height = True,
        additional_inputs_accordion = gr.Accordion(label = '⚙️ Parameters', open = False, render = False),
        additional_inputs = [
            gr.Slider(minimum = 0,
                      maximum = 1, 
                      step = 0.1,
                      value = 0.6, 
                      label = 'Temperature', 
                      render = False),
            gr.Slider(minimum = 128, 
                      maximum = 4096,
                      step = 1,
                      value = 2048, 
                      label = 'Max new tokens', 
                      render = False),
            ],
        examples = [
            ['Saya seorang CPNS di lingkungan Kementerian Hukum dan HAM RI, bekerja dengan sistem (anggota regu pengamanan lapas). Lazimnya PNS dengan jabatan yang sama diberikan libur kerja selama 2 hari. Saya selaku CPNS diwajibkan untuk melaksanakan dinas pada hari libur ke-2 secara berkala tanpa diberikan perintah secara tertulis/surat perintah. Apakah boleh CPNS bekerja dengan jam kerja yang sama dengan PNS? Dan Apakah CPNS boleh bekerja di luar aturan kerja yang sudah ditentukan?'],
            ['Mohon bantu jelaskan apa itu politik luar negeri? Indonesia menganut konsep politik luar negeri apa? Apa dasar hukum politik luar negeri Indonesia?'],
            ['Saya adalah istri kedua dari pernikahan siri, yang mana pernikahannya dilaksanakan tanpa sepengetahuan istri pertama. Dari pernikahan ini kami telah dikaruniai anak. Namun, saya ingin menggugat cerai suami. Bisakah istri siri menggugat cerai atau minta cerai? Bagaimana caranya?'],
            ['Apa yang dimaksud dengan bank perantara dan bank kustodian?Bagaimana langkah untuk meminta eksekusi dari putusan arbitrase nasional? Bisakah kita mengajukan eksekusi terhadap harta salah satu pihak yang tidak disebutkan dalam putusan arbitrase?'],
            ['Adakah hukumnya yang mengatur pembagian persentase/laba dalam mendirikan suatu perusahaan? Dan berapa persenkah yang didapat oleh si pemilik ide untuk mendirikan perusahaan, jika dia tidak menyetor modal sedikit pun atau hanya menjalankan saja?'],
            ],
        cache_examples = False,
    )
    
    gr.Markdown(LICENSE)
    
if __name__ == '__main__':
    demo.launch()