Spaces:
Runtime error
Runtime error
File size: 2,033 Bytes
c4e3c5c c7d81ab 46d5e73 c7d81ab c4e3c5c 46d5e73 a7ab6a0 46d5e73 c7d81ab 34a28ae a7ab6a0 c7d81ab 46d5e73 c7d81ab 46d5e73 c7d81ab 46d5e73 |
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 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import torch
import threading
MODEL = "wmaousley/MiniCrit-1.5B"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
MODEL,
torch_dtype=torch.float16,
device_map="cpu"
)
def generate_stream(prompt):
"""Streaming generator."""
inputs = tokenizer(prompt, return_tensors="pt")
streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
generation_kwargs = dict(
**inputs,
max_new_tokens=200,
temperature=0.7,
do_sample=True,
streamer=streamer
)
thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
for new_token in streamer:
yield new_token
def chat_fn(message, history):
"""Formats chat history + generates streaming reply."""
# Build conversation text
conversation = ""
for user, bot in history:
conversation += f"User: {user}\nMiniCrit: {bot}\n"
conversation += f"User: {message}\nMiniCrit:"
# Stream tokens
reply = ""
for token in generate_stream(conversation):
reply += token
yield reply
# -------- UI --------
with gr.Blocks(theme=gr.themes.Base()) as demo:
gr.Markdown(
"""
<h1 style='text-align:center; color:#00eaff;'>
MiniCrit-1.5B Chat UI 🚀
</h1>
<p style='text-align:center; color:gray;'>Enhanced Streaming Interface</p>
"""
)
chatbox = gr.Chatbot(
label="MiniCrit-1.5B",
height=500
)
with gr.Row():
msg = gr.Textbox(
placeholder="Ask something...",
label="Message",
scale=10
)
send = gr.Button("Send", variant="primary")
clear = gr.Button("Clear")
send.click(chat_fn, [msg, chatbox], chatbox)
send.click(lambda: "", None, msg)
clear.click(lambda: [], None, chatbox)
demo.launch(debug=True)
|