Create app.py

app.py

@@ -0,0 +1,210 @@
+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,
+        'top_p' : 0.95,
+        '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()