Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| import os | |
| from threading import Thread | |
| import spaces | |
| token = os.environ["HF_TOKEN"] | |
| model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it", | |
| # torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, | |
| torch_dtype=torch.float16, | |
| token=token) | |
| tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it",token=token) | |
| # using CUDA for an optimal experience | |
| # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| device = torch.device('cuda') | |
| model = model.to(device) | |
| def chat(message, history): | |
| chat = [] | |
| for item in history: | |
| chat.append({"role": "user", "content": item[0]}) | |
| if item[1] is not None: | |
| chat.append({"role": "assistant", "content": item[1]}) | |
| chat.append({"role": "user", "content": message}) | |
| messages = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True) | |
| # Tokenize the messages string | |
| model_inputs = tokenizer([messages], return_tensors="pt").to(device) | |
| 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.75, | |
| num_beams=1, | |
| ) | |
| t = Thread(target=model.generate, kwargs=generate_kwargs) | |
| t.start() | |
| # Initialize an empty string to store the generated text | |
| partial_text = "" | |
| for new_text in streamer: | |
| # print(new_text) | |
| partial_text += new_text | |
| # Yield an empty string to cleanup the message textbox and the updated conversation history | |
| yield partial_text | |
| demo = gr.ChatInterface(fn=chat, | |
| chatbot=gr.Chatbot(show_label=True, show_share_button=True, show_copy_button=True, likeable=True, layout="bubble", bubble_full_width=False), | |
| theme="upsatwal/mlsc_tiet", | |
| examples=[["Write me a poem about Machine Learning."]], | |
| title="Text Streaming") | |
| demo.launch() |