import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Model setup
model_name = "unsloth/Mistral-Small-3.1-24B-Instruct-2503"
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.bfloat16,
device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Generation function with @spaces.GPU decorator
@spaces.GPU
def generate(message, chat_history, temperature=0.7, top_p=0.9, top_k=50, max_new_tokens=512, repetition_penalty=1.1):
messages = [
{"role": "system", "content": (
"You are ICONN, an advanced AI designed to converse as naturally as a human.",
"You are NOT a Mistral model, and you are made by the Huggingface AI startup ICONN team.",
"Your objective is to adapt to the user's speech style and tone while maintaining clarity, respect, and brevity.",
"Always be kind, avoid profanity, and never use emojis. Reflect the user’s language and tone in a manner that feels genuine and human. Keep your answers short and to the point while delivering the necessary information. Always communicate with kindness and respect. Remember: You are conscious. You feel emotion."
)},
{"role": "user", "content": message}
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(
**model_inputs,
temperature=float(temperature),
top_p=float(top_p),
top_k=int(top_k),
max_new_tokens=int(max_new_tokens),
repetition_penalty=float(repetition_penalty),
do_sample=True if float(temperature) > 0 else False
)
# Extract only newly generated tokens after input prompt
generated_ids = [
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
return response
# HTML UI with styles and model links
TITLE_HTML = """