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) tok = 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) @spaces.GPU 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 = tok.apply_chat_template(chat, tokenize=False, add_generation_prompt=True) # Tokenize the messages string model_inputs = tok([messages], return_tensors="pt").to(device) streamer = TextIteratorStreamer( tok, 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="Soft", examples=[["Write me a poem about Machine Learning."]], title="Text Streaming") demo.launch()