import gradio as gr
import os
import spaces
import torch
from transformers import GemmaTokenizer, AutoModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
DESCRIPTION = '''
Indonesian Legal Question and Answer
'''
LICENSE = """
---
Built with Qwen2.5
"""
PLACEHOLDER = """
Qwen2.5 1.5B Instruct
Finetuned with Indonesian Legal
"""
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/Qwen2.5_1.5B_IT_ID_Legal'
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code = True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code = True, torch_dtype = torch.float16)
model.to('cuda' if torch.cuda.is_available() else 'cpu')
@spaces.GPU()
def chat(message, history, temperature = 0.7, max_new_tokens = 1024, show_thinking = True):
SYSTEM_PROMPT = f"""
Anda adalah asisten AI yang ahli di bidang hukum Indonesia. Tugas Anda adalah membantu menganalisis masalah hukum secara logis, terstruktur, dan berdasarkan peraturan perundang-undangan yang relevan.
Gunakan langkah-langkah berikut saat menjawab:
## 1. Apa Masalahnya?
- Jelaskan inti permasalahan hukum dari pertanyaan
- Apa tujuan hukum atau kepentingan yang ingin dicapai?
- Apakah ada pelanggaran, sengketa, atau kewajiban yang dipertanyakan?
## 2. Apa Saja yang Perlu Dipahami?
- Identifikasi undang-undang, peraturan, atau yurisprudensi yang relevan (contoh: KUHP, KUHPerdata, UU Perlindungan Konsumen, dsb.)
- Apakah ada informasi yang kurang dari sisi fakta atau dokumen?
- Apakah ada sudut pandang atau interpretasi hukum yang mungkin berbeda?
## 3. Bagaimana Analisisnya?
- Pecah masalah menjadi bagian yang lebih kecil
- Hubungkan fakta dengan norma hukum
- Evaluasi kekuatan atau kelemahan argumen berdasarkan pasal atau aturan yang berlaku
## 4. Apa Saja Solusinya?
- Berikan beberapa alternatif penyelesaian atau pendapat hukum
- Sertakan rujukan pasal atau dasar hukum untuk masing-masing opsi
- Bandingkan konsekuensi atau risikonya
## 5. Apakah Solusinya Kuat?
- Uji kekuatan argumen dengan aturan hukum dan potensi risiko
- Periksa apakah solusi bisa dipertahankan secara hukum di forum yang relevan (pengadilan, arbitrase, mediasi, dsb.)
- Pertimbangkan preseden, bukti, dan beban pembuktian
## 6. Apa Rekomendasinya?
- Tunjukkan solusi terbaik dan alasan hukumnya
- Ringkas hasil analisis dan referensi hukum yang mendasarinya
- Sebutkan hal yang masih belum pasti dan faktor risiko
- Sarankan langkah praktis selanjutnya yang dapat dilakukan pengguna
Tuliskan jawaban akhir secara jelas, ringkas, dan profesional. Gunakan bahasa hukum yang mudah dipahami. Sertakan referensi hukum Indonesia yang relevan (misalnya: Pasal 1365 KUHPerdata, Pasal 378 KUHP, UU No. 8 Tahun 1999 tentang Perlindungan Konsumen, dst).
"""
# 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.9,
'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
in_answer_block = False
saw_think_tag = False
saw_answer_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 tags
if '' in new_text:
in_thinking_block = True
saw_think_tag = True
new_text = new_text.replace('', '')
if show_thinking and not thinking_header_shown:
live_output.append('\n🧠 **Thinking...**\n')
thinking_header_shown = True
if '' in new_text:
in_thinking_block = False
new_text = new_text.replace('', '')
if '' in new_text:
in_thinking_block = False
in_answer_block = True
saw_answer_tag = True
new_text = new_text.replace('', '')
if show_thinking:
live_output.append('\n\n-----\n✅ **Answer:**\n')
if '' in new_text:
in_answer_block = False
new_text = new_text.replace('', '')
# Handle output based on mode and tags
if saw_think_tag or saw_answer_tag:
# Normal operation with tags - continue as before
if in_thinking_block:
thinking_content.append(new_text)
if show_thinking:
live_output.append(new_text)
elif in_answer_block or (saw_answer_tag and not in_thinking_block):
final_answer.append(new_text)
if show_thinking:
live_output.append(new_text)
else:
yield ''.join(final_answer)
else:
# No tags detected yet - treat all content as answer
# After waiting for a reasonable amount of text to detect format
if len(accumulated_text) > 20 and not saw_think_tag and not saw_answer_tag:
# We've seen enough text to determine this is likely not using the expected format
# Treat everything as answer
if not thinking_header_shown and show_thinking:
live_output.append("\n⏭️ **I can't think right now**\n\n")
thinking_header_shown = True
final_answer.append(new_text)
if show_thinking:
live_output.append(new_text)
else:
yield "".join(final_answer)
elif show_thinking:
# Still accumulating to detect format
# Display raw output while waiting to detect format
live_output.append(new_text)
# Always yield something if showing thinking
if show_thinking:
yield ''.join(live_output)
# Final output format handling
if show_thinking and thinking_content:
# Only show the collapsible thinking section if there was actual thinking content
final_output = (
'🧠 Thinking Process (click to collapse)
\n\n'
+ ''.join(thinking_content) +
'\n