Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| # Model definitions | |
| PRIMARY_MODEL = "Smilyai-labs/Sam-reason-A1" | |
| FALLBACK_MODEL = "Smilyai-labs/Sam-reason-S2.1" | |
| USAGE_LIMIT = 10 | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Globals for models and tokenizers | |
| primary_model, primary_tokenizer = None, None | |
| fallback_model, fallback_tokenizer = None, None | |
| # IP-based usage tracking | |
| usage_counts = {} | |
| def load_models(): | |
| global primary_model, primary_tokenizer, fallback_model, fallback_tokenizer | |
| primary_tokenizer = AutoTokenizer.from_pretrained(PRIMARY_MODEL) | |
| primary_model = AutoModelForCausalLM.from_pretrained(PRIMARY_MODEL).to(device).eval() | |
| fallback_tokenizer = AutoTokenizer.from_pretrained(FALLBACK_MODEL) | |
| fallback_model = AutoModelForCausalLM.from_pretrained(FALLBACK_MODEL).to(device).eval() | |
| return f"Models loaded: {PRIMARY_MODEL} + fallback {FALLBACK_MODEL}" | |
| def generate_stream(prompt, use_fallback=False, max_length=100, temperature=0.7, top_p=0.9): | |
| model = fallback_model if use_fallback else primary_model | |
| tokenizer = fallback_tokenizer if use_fallback else primary_tokenizer | |
| input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device) | |
| generated = input_ids | |
| output_text = tokenizer.decode(input_ids[0]) | |
| for _ in range(max_length): | |
| outputs = model(generated) | |
| logits = outputs.logits[:, -1, :] / temperature | |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) | |
| probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) | |
| mask = probs > top_p | |
| mask[..., 1:] = mask[..., :-1].clone() | |
| mask[..., 0] = 0 | |
| filtered = logits.clone() | |
| filtered[:, sorted_indices[mask]] = -float("Inf") | |
| next_token = torch.multinomial(torch.softmax(filtered, dim=-1), 1) | |
| generated = torch.cat([generated, next_token], dim=-1) | |
| new_text = tokenizer.decode(next_token[0]) | |
| output_text += new_text | |
| yield output_text | |
| if next_token.item() == tokenizer.eos_token_id: | |
| break | |
| def respond(msg, history, reasoning_enabled, request: gr.Request): | |
| ip = request.client.host if request else "unknown" | |
| usage_counts[ip] = usage_counts.get(ip, 0) + 1 | |
| use_fallback = usage_counts[ip] > USAGE_LIMIT | |
| model_used = "A1" if not use_fallback else "Fallback S2.1" | |
| prefix = "/think " if reasoning_enabled else "/no_think " | |
| prompt = prefix + msg.strip() | |
| history = history + [[msg, ""]] | |
| for output in generate_stream(prompt, use_fallback): | |
| history[-1][1] = output + f" ({model_used})" | |
| yield history, history | |
| def clear_chat(): | |
| return [], [] | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🤖 SmilyAI Reasoning Chat • Token-by-Token + IP Usage Limits") | |
| model_status = gr.Textbox(label="Model Load Status", interactive=False) | |
| chat_box = gr.Chatbot(label="Chat", type="tuples") | |
| chat_state = gr.State([]) | |
| with gr.Row(): | |
| user_input = gr.Textbox(placeholder="Your message here...", show_label=False, scale=6) | |
| reason_toggle = gr.Checkbox(label="Reason", value=True, scale=1) | |
| send_btn = gr.Button("Send", scale=1) | |
| clear_btn = gr.Button("Clear Chat") | |
| model_status.value = load_models() | |
| send_btn.click( | |
| respond, | |
| inputs=[user_input, chat_state, reason_toggle], | |
| outputs=[chat_box, chat_state] | |
| ) | |
| clear_btn.click(fn=clear_chat, inputs=[], outputs=[chat_box, chat_state]) | |
| demo.queue() | |
| demo.launch() | |