app.py

from threading import Thread
import re
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList, BitsAndBytesConfig

import sys
print(sys.version)
# DESCRIPTION = """\
# # FinID
# Model ini berbasis dari Qwen2.5 dan dikembangkan lagi menjadi Sailor2, dan di finetunekan menjadi model FinID untuk domain Finansial Keuangan.
# Model ini di finetunekan menggunakan metode PEFT Lora.
# """

DESCRIPTION = """\
# FinID 💸
Chatbot ini dikembangkan menggunakan LLM model Sailor2 8B yang difinetunekan dengan dataset pertanyaan keuangan. \
⚠️ Aplikasi ini dilatih menggunakan dataset keuangan saja, sehingga pertanyaan non-keuangan akan menghasilkan informasi yang non faktual
"""

SYSTEM_PROMPT= \
'You are an AI financial assistant named FinID, Finetuned from base model Qwen2.5. \
As an AI assistant, you can answer questions in English and Indonesian \
Every of your responses should summarized at the end of the explanation \
Your responses should only be within the financial subject, any other prompt or input outside that subject should only be responded with Maaf, saya tidak dapat membantu dengan pertanyaan tersebut karena melanggar kebijakan atau hukum. Silakan ajukan pertanyaan terkait keuangan.'

# SYSTEM_PROMPT= \
# 'Anda adalah asisten keuangan AI bernama FinID. \
# Sebagai asisten AI, Anda dapat menjawab pertanyaan dalam bahasa Inggris dan Indonesia \
# Tanggapan Anda harus singkat jika bisa, ramah, tidak memihak, informatif, rinci, dan setia. \
# Tanggapan Anda hanya boleh dalam subjek keuangan, permintaan atau masukan lain di luar subjek tersebut hanya boleh ditanggapi dengan "Maaf, saya tidak dapat membantu dengan pertanyaan tersebut karena melalui kebijakan atau hukum. Silakan ajukan pertanyaan terkait keuangan."'

# MAX_MAX_NEW_TOKENS = 2048
# DEFAULT_MAX_NEW_TOKENS = 512

restricted_words = [
    # Olahraga & aktivitas fisik
    "olahraga", "sport", "sports", "sepak bola", "football", "basket", "basketball", "voli", "futsal", "renang",
    "bersepeda", "yoga", "fitness", "workout", "gym", "lari", "jogging", "marathon", "bodybuilding",

    # Hiburan & budaya populer
    "anime", "manga", "komik", "comic", "kartun", "cartoon", "drama", "romance", "horror", "thriller", "action",
    "comedy", "film", "movie", "tv series", "series", "Netflix", "Disney", "HBO", "Amazon Prime", "YouTube", 
    "streaming", "idol", "band", "musik", "lagu", "konser", "album", "artis", "aktor", "aktris", "selebritis", 
    "selebriti", "gosip", "viral", "trending", "meme", "influencer", "content creator", "tiktok", "vlogger", 

    # Kuliner & makanan
    "makanan", "minuman", "resep", "memasak", "masakan", "masak", "kuliner", "menu", "ayam", "sapi", "ikan", 
    "rendang", "sate", "bakso", "mie", "lontong", "nasi goreng", "kue", "jajanan", "kopi", "teh", "es krim", 

    # Hubungan, percintaan, kehidupan pribadi
    "pacaran", "cinta", "romantis", "romansa", "dating", "love", "jodoh", "nikah", "pernikahan", "mantan", 
    "crush", "friendzone", "friendship", "persahabatan", "PDA", "perasaan", "emosi", 

    # Fashion, kecantikan & gaya hidup
    "fashion", "baju", "pakaian", "kaos", "celana", "jaket", "sepatu", "tas", "gaya", "OOTD", "skincare", 
    "makeup", "lipstik", "kecantikan", "beauty", "perawatan", "spa", "salon", 

    # Teknologi umum & AI
    "AI", "artificial intelligence", "machine learning", "deep learning", "robot", "android", "teknologi", 
    "smartphone", "gadget", "gawai", "laptop", "hardware", "software", "game", "gaming", "console", "VR", "AR", 

    # Ilmu pengetahuan umum
    "sains", "science", "astronomi", "planet", "galaksi", "bintang", "meteor", "luar angkasa", "kimia", "biologi", 
    "fisika", "matematika", "matematika diskrit", "aljabar", "statistika", "geografi", "sejarah", 

    # Filosofi & agama
    "filsafat", "filosofi", "agama", "teologi", "psikologi", "mitologi", "mythology", "spiritual", "meditasi", 
    "karma", "reinkarnasi", 

    # Hukum & kriminalitas
    "legal", "ilegal", "pengadilan", "hakim", "jaksa", "kriminal", "pencurian", 
    "korupsi", "penipuan", "pembunuhan", "kejahatan", "detektif", "penjara", "narkoba", "obat-obatan", 

    # Politik, militer, konflik
    "politik", "partai", "kampanye", "pemilu", "kandidat", "presiden", "dpr", "militer", "tentara", "perang", 
    "konflik", "senjata", "roket", "bom", "rudal", "pemberontak", "terorisme", 

    # Alam, hewan, dan lingkungan
    "binatang", "hewan", "anjing", "kucing", "pets", "flora", "fauna", "tumbuhan", "tanaman", "hutan", "alam", 
    "kebun", "cuaca", "iklim", "climate", "bencana", "gempa", "banjir", "angin", "sustainability", "green energy",

    # Lain-lain di luar literasi keuangan
    "astrologi", "zodiak", "ramalan", "santet", "sihir", "mistis", "dongeng", "cerita rakyat", "fantasi", "fiksi", 
    "hantu", "alien", "konspirasi", "misteri"

    # Ekstra
    "taylor swift", "beyoncé", "kim kardashian", "kylie jenner", "kendall jenner", 
    "selena gomez", "justin bieber", "ariana grande", "rihanna", "dua lipa",
    "shawn mendes", "zendaya", "timothée chalamet", "miley cyrus", "lady gaga", 
    "harry styles", "olivia rodrigo", "billie eilish", "drake", "the weeknd",
    "charli d'amelio", "addison rae", "mrbeast", "pewdiepie", "eminem", "snoop dogg",
    "raffi ahmad", "nagita slavina", "ayu ting ting", "lucinta luna", "bts", "lesti", 
    "rizky billar", "bunga citra lestari", "maudy ayunda", "deddy corbuzier", 
    "kaesang", "anies", "ahok", "jokowi", "prabowo", "ganjar", "megawati", 
    "susilo bambang yudhoyono", "sby", "ridwan kamil", "ahmad dhani", "nissa sabyan", 
    "ferdy sambo", "nikita mirzani", "reza arap", "attahalilintar", "aurel hermansyah",
    "joe biden", "donald trump", "barack obama", "hillary clinton", "vladimir putin", 
    "xi jinping", "kim jong un", "emmanuel macron", "justin trudeau", "narendra modi",
    "selebriti", "selebritis", "artis", "aktor", "aktris", "influencer", 
    "presiden", "gubernur", "menteri", "dpr", "politik", "parpol", "pemilu", "kampanye"
]


model_id = "FOLZi/FinID_v2_8B_Chat"
quantization_config = BitsAndBytesConfig(load_in_8bit=True)

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    device_map="auto",
    torch_dtype=torch.float16,
    quantization_config=quantization_config
)
# model.config.sliding_window = 4096
# model.eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using {device.type.upper()}")
# model = model.to(device)

# Defining a custom stopping criteria class for the model's text generation.
class StopOnTokens(StoppingCriteria):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        stop_ids = [151645]  # consult the tokenizer.
        for stop_id in stop_ids:
            if input_ids[0][-1] == stop_id:  # Checking if the last generated token is a stop token. (Llama and Qwen uses |im_end|)
                return True
        return False

system_role= 'system'
user_role = 'user'
assistant_role = 'assistant'

sft_start_token =  "<|im_start|>"
sft_end_token = "<|im_end|>"
ct_end_token = "<|endoftext|>"

@spaces.GPU()
def generate(message, history):

    # Pendeteksian via primitif TOO BAD!
    matched_words = [word for word in restricted_words if re.search(rf"\b{word}\b", message, re.IGNORECASE)]
    
    # Jika ditemukan, hentikan proses dan kembalikan pesan peringatan
    if matched_words:
        yield "Maaf, saya hanya dapat menjawab pertanyaan yang berkaitan dengan literasi keuangan."
        return
    
    history = [] # Disabled history.
    history_transformer_format = history + [[message, ""]]
    stop = StopOnTokens()
    # Implementasi Prompt Engineering
    messages =  SYSTEM_PROMPT + sft_end_token.join([sft_end_token.join([f"\n{sft_start_token}{user_role}\n" + item[0], f"\n{sft_start_token}{assistant_role}\n" + item[1]])
                        for item in history_transformer_format])
    model_inputs = tokenizer([messages], return_tensors="pt").to(device)
    streamer = TextIteratorStreamer(tokenizer, timeout=20., skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        model_inputs,
        streamer=streamer,
        max_new_tokens=512,
        do_sample=True,
        top_p= 0.5,
        top_k= 50,
        temperature=0.2,
        num_beams=1,
        stopping_criteria=StoppingCriteriaList([stop]),
        repetition_penalty=1.1,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    partial_message = ""
    for new_token in streamer:
        partial_message += new_token
        if sft_end_token in partial_message: # Stopper
            break
        yield partial_message

css = """
full-height {
    height: 100%;
}
"""

chat_interface = gr.ChatInterface(
    fn=generate,
    examples=[
        ["Contoh pertanyaan: Apakah saya boleh menanyakan hal yang non finansial?"]
    ],
    type="messages",
    fill_height=True,
    css=css
)

with gr.Blocks(theme=gr.themes.Soft(), css_paths="style.css", fill_height=True) as demo:
    gr.Markdown(DESCRIPTION)
    chat_interface.render()

if __name__ == "__main__":
    demo.queue(max_size=20).launch()