Spaces:
Sleeping
Sleeping
File size: 3,297 Bytes
7bd6e67 a7cccd7 7bd6e67 a6ab0c4 1cc47b4 76a6d8e 79a49f5 205ebd7 1cc47b4 4759f07 e4cd356 4759f07 e4cd356 7b061fd 205ebd7 1cc47b4 4759f07 e4cd356 205ebd7 67fdcc7 205ebd7 7bd6e67 7b061fd 7bd6e67 e8081aa 7bd6e67 e4cd356 76a6d8e 7bd6e67 e4cd356 7bd6e67 e4cd356 7bd6e67 e4cd356 7bd6e67 7b061fd df914b7 7bd6e67 84facec 06dfd60 7bd6e67 df914b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
import os
from threading import Thread
from typing import Iterator
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
ACCESS_TOKEN = os.getenv("HF_TOKEN", "")
model_id = "mistralai/Mistral-Nemo-Instruct-2407"
#filename = "Mistral-Nemo-Instruct-2407-Q6_K_L.gguf"
model = AutoModelForCausalLM.from_pretrained(
model_id,
#gguf_file=filename,
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=False,
token=ACCESS_TOKEN)
tokenizer = AutoTokenizer.from_pretrained(
model_id,
#gguf_file=filename,
trust_remote_code=False,
token=ACCESS_TOKEN)
tokenizer.use_default_system_prompt = False
model.config.gradient_checkpointing = True
@spaces.GPU
def generate(
message: str,
system_prompt: str,
max_new_tokens: int = 1024,
temperature: float = 0.01,
top_p: float = 0.7,
) -> Iterator[str]:
conversation = []
if system_prompt:
conversation.append({"role": "system", "content": system_prompt})
conversation.append({"role": "user", "content": message})
input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
input_ids = input_ids.to(model.device)
'''
terminators = [
tokenizer.eos_token_id,
tokenizer.convert_tokens_to_ids("<|eot_id|>")
]
'''
streamer = TextIteratorStreamer(tokenizer, timeout=600.0, skip_prompt=True, skip_special_tokens=True)
generate_kwargs = dict(
{"input_ids": input_ids},
streamer=streamer,
max_new_tokens=max_new_tokens,
#eos_token_id=terminators,
do_sample=True,
top_p=top_p,
temperature=temperature,
num_beams=1,
pad_token_id=tokenizer.eos_token_id,
)
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
outputs = []
for text in streamer:
outputs.append(text)
yield "".join(outputs)
chat_interface = gr.Interface(
fn=generate,
inputs=[
gr.Textbox(lines=2, placeholder="Prompt", label="Prompt"),
],
outputs="text",
additional_inputs=[
gr.Textbox(label="System prompt", lines=6),
gr.Slider(
label="Max new tokens",
minimum=1,
maximum=MAX_MAX_NEW_TOKENS,
step=1,
value=DEFAULT_MAX_NEW_TOKENS,
),
gr.Slider(
label="Temperature",
minimum=0.1,
maximum=4.0,
step=0.01,
value=0.01,
),
gr.Slider(
label="Top-p (nucleus sampling)",
minimum=0.05,
maximum=1.0,
step=0.01,
value=0.7,
),
],
title="Model testing - mistralai/Mistral-Nemo-Instruct-2407",
description="Provide system settings and a prompt to interact with the model.",
)
chat_interface.queue(max_size=20).launch()
|