import gradio as gr
import os
import torch
from transformers import GemmaTokenizer, AutoModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
DESCRIPTION = '''
Indonesian Legal Assistant
'''
LICENSE = """
---
Built with Deepseek Qwen3
"""
PLACEHOLDER = """
DeepSeek-R1-0528-Qwen3-8B
Finetuned with Indonesian Legal Consulting Data
"""
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:
...
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 '' in new_text:
in_thinking_block = True
saw_think_tag = True
new_text = new_text.replace('', '')
if not thinking_header_shown:
live_output.append('\n🧠 **Sedang berfikir...**\n')
thinking_header_shown = True
if '' in new_text:
in_thinking_block = False
new_text = new_text.replace('', '')
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 = (
'🧠 Proses berfikir (klik untuk melihat)
\n\n'
+ ''.join(thinking_content) +
'\n