Spaces:
Sleeping
Sleeping
File size: 3,614 Bytes
36942d4 852d26e cee13f4 f0687e5 341bd22 852d26e fc2aea6 f0687e5 852d26e 341bd22 f0687e5 341bd22 f0687e5 852d26e 485360d 852d26e 2918965 852d26e 341bd22 b56c1a4 852d26e 341bd22 f0687e5 63d4a2a f0687e5 63d4a2a 341bd22 e2b2cdb f0687e5 |
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 118 119 |
import os
import threading
import gradio as gr
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TextIteratorStreamer,
)
# Define your models
MODEL_PATHS = {
"LeCarnet-3M": "MaxLSB/LeCarnet-3M",
"LeCarnet-8M": "MaxLSB/LeCarnet-8M",
"LeCarnet-21M": "MaxLSB/LeCarnet-21M",
}
# Add your Hugging Face token
hf_token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
if not hf_token:
raise ValueError("HUGGINGFACEHUB_API_TOKEN environment variable not set.")
# Load tokenizers & models - only load one initially
tokenizer = None
model = None
def load_model(model_name):
"""Loads the specified model and tokenizer."""
global tokenizer, model
if model_name not in MODEL_PATHS:
raise ValueError(f"Unknown model: {model_name}")
print(f"Loading {model_name}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATHS[model_name], token=hf_token)
model = AutoModelForCausalLM.from_pretrained(MODEL_PATHS[model_name], token=hf_token)
model.eval()
print(f"{model_name} loaded.")
# Initial model load
initial_model = list(MODEL_PATHS.keys())[0]
load_model(initial_model)
def respond(
prompt: str,
chat_history,
model_choice: str,
max_tokens: int,
temperature: float,
top_p: float,
):
global tokenizer, model
# Reload model if it's not the currently loaded one
if model.config._name_or_path != MODEL_PATHS[model_choice]:
load_model(model_choice)
inputs = tokenizer(prompt, return_tensors="pt")
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt=False,
skip_special_tokens=True,
)
generate_kwargs = dict(
**inputs,
streamer=streamer,
max_new_tokens=max_tokens,
do_sample=True,
temperature=temperature,
top_p=top_p,
eos_token_id=tokenizer.eos_token_id,
)
thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
thread.start()
accumulated = ""
for new_text in streamer:
accumulated += new_text
yield accumulated
with gr.Blocks(css=css, fill_width=True) as demo:
with gr.Row():
with gr.Column(scale=1):
model_dropdown = gr.Dropdown(
choices=list(MODEL_PATHS.keys()),
value=initial_model,
label="Choose Model",
interactive=True
)
max_tokens_slider = gr.Slider(
1, 512, value=512, step=1, label="Max new tokens"
)
temperature_slider = gr.Slider(
0.1, 2.0, value=0.7, step=0.1, label="Temperature"
)
top_p_slider = gr.Slider(
0.1, 1.0, value=0.9, step=0.05, label="Top‑p"
)
with gr.Column(scale=3):
chatbot = gr.ChatInterface(
fn=respond,
additional_inputs=[
model_dropdown,
max_tokens_slider,
temperature_slider,
top_p_slider,
],
examples=[
["Il était une fois un petit garçon qui vivait dans un village paisible."],
["Il était une fois une grenouille qui rêvait de toucher les étoiles chaque nuit depuis son étang."],
["Il était une fois un petit lapin perdu"],
],
cache_examples=False,
submit_btn="Generate",
avatar_images=(None, "media/le-carnet.png")
)
if __name__ == "__main__":
demo.queue()
demo.launch() |