FinMax

Paused

FinMax / app.py

Update app.py

050c155 verified about 2 years ago

10.3 kB

import os
import json
import subprocess
from threading import Thread

import torch
import spaces
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
from sentence_transformers import SentenceTransformer
from datasets import load_dataset


subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

MODEL_ID = os.environ.get("MODEL_ID")
CHAT_TEMPLATE = os.environ.get("CHAT_TEMPLATE")
MODEL_NAME = MODEL_ID.split("/")[-1]
CONTEXT_LENGTH = int(os.environ.get("CONTEXT_LENGTH"))
COLOR = os.environ.get("COLOR")
EMOJI = os.environ.get("EMOJI")
DESCRIPTION = os.environ.get("DESCRIPTION")
token = os.environ["HF_TOKEN"]

# ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# dataset = load_dataset("gagan3012/edgar-corpus-embeddings")

# data = dataset["train"]
# data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset

def search(query: str, k: int = 3 ):
    """a function that embeds a new query and returns the most probable results"""
    embedded_query = ST.encode(query) # embed new query
    scores, retrieved_examples = data.get_nearest_examples( # retrieve results
        "embeddings", embedded_query, # compare our new embedded query with the dataset embeddings
        k=k # get only top k results
    )
    return scores, retrieved_examples

def format_prompt(prompt,retrieved_documents,k):
    """using the retrieved documents we will prompt the model to generate our responses"""
    PROMPT = f"Question:{prompt}\nContext:"
    for idx in range(k) :
        PROMPT+= f"{retrieved_documents['text'][idx]}\n"
    return PROMPT


# @spaces.GPU()
# def predict(message, history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p):
#     # Format history with a given chat template
#     system_prompt = """You are FinTral, a financial expert specializing in the nuanced analysis
# of financial statements and a wide array of data-driven financial
# tasks. For each prompt you are given, think step by step. Sometimes, you must extract relevant information to proceed with the
# problem.
# Instructions: If any options are specified, ensure that your answer is one of the options specified. Do not explain why you think the answer is correct.
# """
#     if CHAT_TEMPLATE == "ChatML":
#         stop_tokens = ["<|endoftext|>", "<|im_end|>", "<|end_of_text|>", "<|eot_id|>", "assistant"]
#         instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'
#         # scores , retrieved_documents = search(message, k=2)
#         # final_message = format_prompt(message,retrieved_documents,k=2)
#         for human, assistant in history:
#             instruction += '<|im_start|>user\n' + human + '\n<|im_end|>\n<|im_start|>assistant\n' + assistant
#         instruction += '\n<|im_start|>user\n' + message + '\n<|im_end|>\n<|im_start|>assistant\n'
#     elif CHAT_TEMPLATE == "Mistral Instruct":
#         stop_tokens = ["</s>", "[INST]", "[INST] ", "<s>", "[/INST]", "[/INST] "]
#         instruction = '<s>[INST] ' + system_prompt
#         for human, assistant in history:
#             instruction += human + ' [/INST] ' + assistant + '</s>[INST]'
#         instruction += ' ' + message + ' [/INST]'
#     else:
#         messages = [{"role":"system","content":system_prompt},{"role":"user","content":message}]
#         for human, assistant in history:
#             messages.append({"role":"user","content":human},{"role":"assistant","content":assistant})
#     # tell the model to generate
#         instruction = tokenizer.apply_chat_template(
#             messages,
#             add_generation_prompt=True,
#             tokenize=False
#         )
#     print(instruction)
    
#     streamer = TextIteratorStreamer(tokenizer, timeout=90.0, skip_prompt=True, skip_special_tokens=True)
#     enc = tokenizer([instruction], return_tensors="pt", padding=True, truncation=True, max_length=CONTEXT_LENGTH)
#     input_ids, attention_mask = enc.input_ids, enc.attention_mask

#     if input_ids.shape[1] > CONTEXT_LENGTH:
#         input_ids = input_ids[:, -CONTEXT_LENGTH:]

#     generate_kwargs = dict(
#         {"input_ids": input_ids.to(device), "attention_mask": attention_mask.to(device)},
#         streamer=streamer,
#         do_sample=True,
#         temperature=temperature,
#         max_new_tokens=max_new_tokens,
#         top_k=top_k,
#         repetition_penalty=repetition_penalty,
#         top_p=top_p
#     )
#     t = Thread(target=model.generate, kwargs=generate_kwargs)
#     t.start()
#     outputs = []
#     for new_token in streamer:
#         outputs.append(new_token)
#         if new_token in stop_tokens:
#             break
#         yield "".join(outputs)


@spaces.GPU()
def stream_chat(
    message: str, 
    history: list,
    system_prompt: str,
    temperature: float = 0.8, 
    max_new_tokens: int = 1024, 
    top_p: float = 1.0, 
    top_k: int = 20, 
    penalty: float = 1.2,
):
    print(f'message: {message}')
    print(f'history: {history}')
    system = """You are FinTral, a financial expert specializing in the nuanced analysis
of financial statements and a wide array of data-driven financial
tasks. For each prompt you are given, think step by step. Sometimes, you must extract relevant information to proceed with the
problem.
Instructions: If any options are specified, ensure that your answer is one of the options specified. Do not explain why you think the answer is correct.
"""

    conversation = [
        {"role": "system", "content": system}
    ]
    for prompt, answer in history:
        conversation.extend([
            {"role": "user", "content": prompt}, 
            {"role": "assistant", "content": answer},
        ])

    conversation.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
    
    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=input_ids, 
        max_new_tokens = max_new_tokens,
        do_sample = False if temperature == 0 else True,
        top_p = top_p,
        top_k = top_k,
        temperature = temperature,
        repetition_penalty=1.2,
        eos_token_id=[128001,128008,128009],
        streamer=streamer,
    )

    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=generate_kwargs)
        thread.start()
        
    buffer = ""
    for new_text in streamer:
        buffer += new_text
        yield buffer

# Load model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type= "nf4"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID,token=token)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    quantization_config=quantization_config,
    attn_implementation="flash_attention_2",
    token=token
)

css = """
.message-row {
    justify-content: space-evenly !important;
}
.message-bubble-border {
    border-radius: 6px !important;
}
.message-buttons-bot, .message-buttons-user {
    right: 10px !important;
    left: auto !important;
    bottom: 2px !important;
}
.dark.message-bubble-border {
    border-color: #15172c !important;
}
.dark.user {
    background: #10132c !important;
}
.dark.assistant.dark, .dark.pending.dark {
    background: #020417 !important;
}
"""

TITLE = "<h1><center>FinTral</center></h1>"

PLACEHOLDER = """
<center>
<p>Hi! How can I help you today?</p>
</center>
"""


CSS = """
.duplicate-button {
    margin: auto !important;
    color: white !important;
    background: black !important;
    border-radius: 100vh !important;
}
h3 {
    text-align: center;
}
"""

# Create Gradio interface
chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)

with gr.Blocks(css=CSS, theme="soft") as demo:
    gr.HTML(TITLE)
    gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
    gr.ChatInterface(
        fn=stream_chat,
        chatbot=chatbot,
        fill_height=True,
        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
        additional_inputs=[
            gr.Textbox(
                value="You are a helpful assistant",
                label="System Prompt",
                render=False,
            ),
            gr.Slider(
                minimum=0,
                maximum=1,
                step=0.1,
                value=0.8,
                label="Temperature",
                render=False,
            ),
            gr.Slider(
                minimum=128,
                maximum=8192,
                step=1,
                value=1024,
                label="Max new tokens",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.1,
                value=1.0,
                label="top_p",
                render=False,
            ),
            gr.Slider(
                minimum=1,
                maximum=20,
                step=1,
                value=20,
                label="top_k",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=2.0,
                step=0.1,
                value=1.2,
                label="Repetition penalty",
                render=False,
            ),
        ],
        examples=[
            ["Help me study vocabulary: write a sentence for me to fill in the blank, and I'll try to pick the correct option."],
            ["What are 5 creative things I could do with my kids' art? I don't want to throw them away, but it's also so much clutter."],
            ["Tell me a random fun fact about the Roman Empire."],
            ["Show me a code snippet of a website's sticky header in CSS and JavaScript."],
        ],
        cache_examples=False,
    )


if __name__ == "__main__":
    demo.launch()

	import os
	import json
	import subprocess
	from threading import Thread

	import torch
	import spaces
	import gradio as gr
	from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
	from sentence_transformers import SentenceTransformer
	from datasets import load_dataset


	subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

	MODEL_ID = os.environ.get("MODEL_ID")
	CHAT_TEMPLATE = os.environ.get("CHAT_TEMPLATE")
	MODEL_NAME = MODEL_ID.split("/")[-1]
	CONTEXT_LENGTH = int(os.environ.get("CONTEXT_LENGTH"))
	COLOR = os.environ.get("COLOR")
	EMOJI = os.environ.get("EMOJI")
	DESCRIPTION = os.environ.get("DESCRIPTION")
	token = os.environ["HF_TOKEN"]

	# ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

	# dataset = load_dataset("gagan3012/edgar-corpus-embeddings")

	# data = dataset["train"]
	# data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset

	def search(query: str, k: int = 3 ):
	"""a function that embeds a new query and returns the most probable results"""
	embedded_query = ST.encode(query) # embed new query
	scores, retrieved_examples = data.get_nearest_examples( # retrieve results
	"embeddings", embedded_query, # compare our new embedded query with the dataset embeddings
	k=k # get only top k results
	)
	return scores, retrieved_examples

	def format_prompt(prompt,retrieved_documents,k):
	"""using the retrieved documents we will prompt the model to generate our responses"""
	PROMPT = f"Question:{prompt}\nContext:"
	for idx in range(k) :
	PROMPT+= f"{retrieved_documents['text'][idx]}\n"
	return PROMPT


	# @spaces.GPU()
	# def predict(message, history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p):
	# # Format history with a given chat template
	# system_prompt = """You are FinTral, a financial expert specializing in the nuanced analysis
	# of financial statements and a wide array of data-driven financial
	# tasks. For each prompt you are given, think step by step. Sometimes, you must extract relevant information to proceed with the
	# problem.
	# Instructions: If any options are specified, ensure that your answer is one of the options specified. Do not explain why you think the answer is correct.
	# """
	# if CHAT_TEMPLATE == "ChatML":
	# stop_tokens = ["<\|endoftext\|>", "<\|im_end\|>", "<\|end_of_text\|>", "<\|eot_id\|>", "assistant"]
	# instruction = '<\|im_start\|>system\n' + system_prompt + '\n<\|im_end\|>\n'
	# # scores , retrieved_documents = search(message, k=2)
	# # final_message = format_prompt(message,retrieved_documents,k=2)
	# for human, assistant in history:
	# instruction += '<\|im_start\|>user\n' + human + '\n<\|im_end\|>\n<\|im_start\|>assistant\n' + assistant
	# instruction += '\n<\|im_start\|>user\n' + message + '\n<\|im_end\|>\n<\|im_start\|>assistant\n'
	# elif CHAT_TEMPLATE == "Mistral Instruct":
	# stop_tokens = ["</s>", "[INST]", "[INST] ", "<s>", "[/INST]", "[/INST] "]
	# instruction = '<s>[INST] ' + system_prompt
	# for human, assistant in history:
	# instruction += human + ' [/INST] ' + assistant + '</s>[INST]'
	# instruction += ' ' + message + ' [/INST]'
	# else:
	# messages = [{"role":"system","content":system_prompt},{"role":"user","content":message}]
	# for human, assistant in history:
	# messages.append({"role":"user","content":human},{"role":"assistant","content":assistant})
	# # tell the model to generate
	# instruction = tokenizer.apply_chat_template(
	# messages,
	# add_generation_prompt=True,
	# tokenize=False
	# )
	# print(instruction)

	# streamer = TextIteratorStreamer(tokenizer, timeout=90.0, skip_prompt=True, skip_special_tokens=True)
	# enc = tokenizer([instruction], return_tensors="pt", padding=True, truncation=True, max_length=CONTEXT_LENGTH)
	# input_ids, attention_mask = enc.input_ids, enc.attention_mask

	# if input_ids.shape[1] > CONTEXT_LENGTH:
	# input_ids = input_ids[:, -CONTEXT_LENGTH:]

	# generate_kwargs = dict(
	# {"input_ids": input_ids.to(device), "attention_mask": attention_mask.to(device)},
	# streamer=streamer,
	# do_sample=True,
	# temperature=temperature,
	# max_new_tokens=max_new_tokens,
	# top_k=top_k,
	# repetition_penalty=repetition_penalty,
	# top_p=top_p
	# )
	# t = Thread(target=model.generate, kwargs=generate_kwargs)
	# t.start()
	# outputs = []
	# for new_token in streamer:
	# outputs.append(new_token)
	# if new_token in stop_tokens:
	# break
	# yield "".join(outputs)


	@spaces.GPU()
	def stream_chat(
	message: str,
	history: list,
	system_prompt: str,
	temperature: float = 0.8,
	max_new_tokens: int = 1024,
	top_p: float = 1.0,
	top_k: int = 20,
	penalty: float = 1.2,
	):
	print(f'message: {message}')
	print(f'history: {history}')
	system = """You are FinTral, a financial expert specializing in the nuanced analysis
	of financial statements and a wide array of data-driven financial
	tasks. For each prompt you are given, think step by step. Sometimes, you must extract relevant information to proceed with the
	problem.
	Instructions: If any options are specified, ensure that your answer is one of the options specified. Do not explain why you think the answer is correct.
	"""

	conversation = [
	{"role": "system", "content": system}
	]
	for prompt, answer in history:
	conversation.extend([
	{"role": "user", "content": prompt},
	{"role": "assistant", "content": answer},
	])

	conversation.append({"role": "user", "content": message})

	input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)

	streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)

	generate_kwargs = dict(
	input_ids=input_ids,
	max_new_tokens = max_new_tokens,
	do_sample = False if temperature == 0 else True,
	top_p = top_p,
	top_k = top_k,
	temperature = temperature,
	repetition_penalty=1.2,
	eos_token_id=[128001,128008,128009],
	streamer=streamer,
	)

	with torch.no_grad():
	thread = Thread(target=model.generate, kwargs=generate_kwargs)
	thread.start()

	buffer = ""
	for new_text in streamer:
	buffer += new_text
	yield buffer

	# Load model
	device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
	quantization_config = BitsAndBytesConfig(
	load_in_4bit=True,
	bnb_4bit_compute_dtype=torch.bfloat16,
	bnb_4bit_use_double_quant=True,
	bnb_4bit_quant_type= "nf4"
	)
	tokenizer = AutoTokenizer.from_pretrained(MODEL_ID,token=token)
	model = AutoModelForCausalLM.from_pretrained(
	MODEL_ID,
	device_map="auto",
	quantization_config=quantization_config,
	attn_implementation="flash_attention_2",
	token=token
	)

	css = """
	.message-row {
	justify-content: space-evenly !important;
	}
	.message-bubble-border {
	border-radius: 6px !important;
	}
	.message-buttons-bot, .message-buttons-user {
	right: 10px !important;
	left: auto !important;
	bottom: 2px !important;
	}
	.dark.message-bubble-border {
	border-color: #15172c !important;
	}
	.dark.user {
	background: #10132c !important;
	}
	.dark.assistant.dark, .dark.pending.dark {
	background: #020417 !important;
	}
	"""

	TITLE = "<h1><center>FinTral</center></h1>"

	PLACEHOLDER = """
	<center>
	<p>Hi! How can I help you today?</p>
	</center>
	"""


	CSS = """
	.duplicate-button {
	margin: auto !important;
	color: white !important;
	background: black !important;
	border-radius: 100vh !important;
	}
	h3 {
	text-align: center;
	}
	"""

	# Create Gradio interface
	chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)

	with gr.Blocks(css=CSS, theme="soft") as demo:
	gr.HTML(TITLE)
	gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
	gr.ChatInterface(
	fn=stream_chat,
	chatbot=chatbot,
	fill_height=True,
	additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
	additional_inputs=[
	gr.Textbox(
	value="You are a helpful assistant",
	label="System Prompt",
	render=False,
	),
	gr.Slider(
	minimum=0,
	maximum=1,
	step=0.1,
	value=0.8,
	label="Temperature",
	render=False,
	),
	gr.Slider(
	minimum=128,
	maximum=8192,
	step=1,
	value=1024,
	label="Max new tokens",
	render=False,
	),
	gr.Slider(
	minimum=0.0,
	maximum=1.0,
	step=0.1,
	value=1.0,
	label="top_p",
	render=False,
	),
	gr.Slider(
	minimum=1,
	maximum=20,
	step=1,
	value=20,
	label="top_k",
	render=False,
	),
	gr.Slider(
	minimum=0.0,
	maximum=2.0,
	step=0.1,
	value=1.2,
	label="Repetition penalty",
	render=False,
	),
	],
	examples=[
	["Help me study vocabulary: write a sentence for me to fill in the blank, and I'll try to pick the correct option."],
	["What are 5 creative things I could do with my kids' art? I don't want to throw them away, but it's also so much clutter."],
	["Tell me a random fun fact about the Roman Empire."],
	["Show me a code snippet of a website's sticky header in CSS and JavaScript."],
	],
	cache_examples=False,
	)


	if __name__ == "__main__":
	demo.launch()