import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer import time import numpy as np from torch.nn import functional as F import os from threading import Thread token = os.environ["HF_TOKEN"] model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,token=token) tok = AutoTokenizer.from_pretrained("google/gemma-2b-it",token=token) # using CUDA for an optimal experience device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) start_message = "" def user(message, history): # Append the user's message to the conversation history return "", history + [[message, ""]] 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, examples=[["Write me a poem about Machine Learning."]], title="gemma 2b-it") demo.launch()