Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| from threading import Thread | |
| print(f"Starting to load the model to memory") | |
| m = AutoModelForCausalLM.from_pretrained( | |
| "stabilityai/stablelm-2-zephyr-1_6b", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, trust_remote_code=True) | |
| tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-zephyr-1_6b", trust_remote_code=True) | |
| # using CUDA for an optimal experience | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| m = m.to(device) | |
| print(f"Sucessfully loaded the model to the memory") | |
| 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=m.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=["hello"], title="GPT Chatboat by Raushan Sharma") | |
| demo.launch(share=True) | |