Spaces:
Sleeping
Sleeping
| """ | |
| This is a Hugging Face Spaces demo for Fin-RWKV-1B5 attention free finanacial export modal. | |
| Author: Umut (Hope) YILDIRIM <[email protected]> | |
| """ | |
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer | |
| from threading import Thread | |
| import torch | |
| tokenizer = AutoTokenizer.from_pretrained("umuthopeyildirim/fin-rwkv-1b5") | |
| model = AutoModelForCausalLM.from_pretrained("umuthopeyildirim/fin-rwkv-1b5") | |
| class StopOnTokens(StoppingCriteria): | |
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: | |
| stop_ids = [29, 0] | |
| for stop_id in stop_ids: | |
| if input_ids[0][-1] == stop_id: | |
| return True | |
| return False | |
| def predict(message, history): | |
| history_transformer_format = history + [[message, ""]] | |
| stop = StopOnTokens() | |
| messages = "".join(["".join(["\nuser :"+item[0], "\nbot:"+item[1]]) # curr_system_message + | |
| for item in history_transformer_format]) | |
| print(messages) | |
| model_inputs = tokenizer([messages], return_tensors="pt") | |
| streamer = TextIteratorStreamer( | |
| tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True) | |
| generate_kwargs = dict( | |
| model_inputs, | |
| streamer=streamer, | |
| max_new_tokens=1024, | |
| do_sample=True, | |
| top_p=0.95, | |
| top_k=1000, | |
| temperature=0.5, | |
| num_beams=1, | |
| stopping_criteria=StoppingCriteriaList([stop]) | |
| ) | |
| t = Thread(target=model.generate, kwargs=generate_kwargs) | |
| t.start() | |
| partial_message = "" | |
| for new_token in streamer: | |
| if new_token != '<': | |
| partial_message += new_token | |
| yield partial_message | |
| def generate_text(prompt, tokenizer, model): | |
| # Tokenize the input | |
| input_ids = tokenizer.encode(prompt, return_tensors="pt") | |
| # Generate a response | |
| output = model.generate(input_ids, max_length=333, num_return_sequences=1) | |
| # Decode the output | |
| generated_text = tokenizer.decode(output[0], skip_special_tokens=True) | |
| return generated_text | |
| title = "# Fin-RWKV: Attention Free Financal Expert (WIP)" | |
| description = """Demo for **Fin-RWKV: Attention Free Financal Expert (WIP)**. | |
| To download the model, please visit [Fin-RWKV: Attention Free Financal Expert (WIP)](https://huggingface.co/umutyildirim/fin-rwkv-1b5).""" | |
| css = """ | |
| #img-display-container { | |
| max-height: 100vh; | |
| } | |
| #img-display-input { | |
| max-height: 80vh; | |
| } | |
| #img-display-output { | |
| max-height: 80vh; | |
| } | |
| """ | |
| with gr.Blocks(css=css) as demo: | |
| gr.Markdown(title) | |
| gr.Markdown(description) | |
| with gr.Tab("Chatbot"): | |
| gr.ChatInterface(predict) | |
| with gr.Tab("E-Commerce"): | |
| gr.Markdown("e-commerce") | |
| with gr.Tab("OpenBB"): | |
| gr.Markdown("openbb") | |
| if __name__ == '__main__': | |
| demo.queue().launch() | |