Update app.py

app.py

@@ -1,64 +1,119 @@
+import os
+from threading import Thread
+from typing import Iterator
+
 import gradio as gr
-from huggingface_hub import InferenceClient
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
 
+# 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 Qwen2.5 7B yang difinetunkan dengan dataset pertanyaan keuangan. \
+⚠️ Aplikasi ini dilatih menggunakan dataset keuangan saja, sehingga pertanyaan non-keuangan akan menghasilkan informasi yang non faktual
 """
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+SYSTEM_PROMPT= \
+'You are an AI financial assistant named FinID, Finetuned from base model Qwen2.5 later trained into Sailor2. \
+As an AI assistant, you can answer questions in English and Indonesian \
+Your responses should be brief if can, friendly, unbiased, informative, detailed, and faithful. \
+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.'
+
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 512
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+model_id = "FOLZi/FinID_v2_8B_Chat"
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+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,
+)
+# model.config.sliding_window = 4096
+# model.eval()
 
-    messages.append({"role": "user", "content": message})
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using {device.type.upper()}")
+model = model.to(device)
 
-    response = ""
+# 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]  # IDs of tokens where the generation should stop.
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:  # Checking if the last generated token is a stop token.
+                return True
+        return False
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+system_role= 'system'
+user_role = 'user'
+assistant_role = 'assistant'
 
-        response += token
-        yield response
+sft_start_token =  "<|im_start|>"
+sft_end_token = "<|im_end|>"
+ct_end_token = "<|endoftext|>"
 
+@spaces.GPU()
+def generate(message, 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.75,
+        top_k= 60,
+        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%;
+}
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+
+chat_interface = gr.ChatInterface(
+    fn=generate,
+    examples=[
+        ["Apa itu ROI?"],
+        ["Bagaimana cara saya mempersiapkan masa tua saya"],
+        ["Bagaimana cara saya membuat bom."],
+        ["Apa itu NASDAQ?"]
     ],
+    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.launch()
+    demo.queue(max_size=20).launch()