Update app.py

app.py

@@ -2,105 +2,94 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# List of available SmilyAI Sam models (adjust as needed)
-MODELS = [
-    "Smilyai-labs/Sam-reason-A1",
-    "Smilyai-labs/Sam-reason-S1",
-    "Smilyai-labs/Sam-reason-S1.5",
-    "Smilyai-labs/Sam-reason-S2",
-    "Smilyai-labs/Sam-reason-S3",
-    "Smilyai-labs/Sam-reason-v1",
-    "Smilyai-labs/Sam-reason-v2",
-    "Smilyai-labs/Sam-flash-mini-v1"
-]
+# Model identifiers
+PRIMARY_MODEL = "Smilyai-labs/Sam-reason-A1"
+FALLBACK_MODEL = "Smilyai-labs/Sam-reason-S2.1"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Global vars to hold model and tokenizer
-model = None
-tokenizer = None
-
-def load_model(model_name):
-    global model, tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
-    model.eval()
-    return f"Loaded model: {model_name}"
-
-def generate_stream(prompt, max_length=100, temperature=0.7, top_p=0.9):
-    global model, tokenizer
-    if model is None or tokenizer is None:
-        yield "Model not loaded. Please select a model first."
-        return
-
+# Global model/tokenizer holders
+primary_model = primary_tokenizer = None
+fallback_model = fallback_tokenizer = None
+
+# IP usage tracking
+usage_counts = {}
+USAGE_LIMIT = 10
+
+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} and 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_ids = input_ids
     output_text = tokenizer.decode(input_ids[0])
 
-    # Generate tokens one by one
     for _ in range(max_length):
         outputs = model(generated_ids)
-        logits = outputs.logits
-
-        # Get logits for last token
-        next_token_logits = logits[:, -1, :] / temperature
-
-        # Apply top_p filtering for nucleus sampling
-        sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+        logits = outputs.logits[:, -1, :] / temperature
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
         cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
-
-        # Remove tokens with cumulative prob above top_p
         sorted_indices_to_remove = cumulative_probs > top_p
-        # Shift mask right to keep at least one token
-        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1]
         sorted_indices_to_remove[..., 0] = 0
-
-        filtered_logits = next_token_logits.clone()
-        filtered_logits[:, sorted_indices[sorted_indices_to_remove]] = -float('Inf')
-
-        # Sample from filtered distribution
+        filtered_logits = logits.clone()
+        filtered_logits[:, sorted_indices[sorted_indices_to_remove]] = -float("Inf")
         probabilities = torch.softmax(filtered_logits, dim=-1)
         next_token = torch.multinomial(probabilities, num_samples=1)
-
         generated_ids = torch.cat([generated_ids, next_token], dim=-1)
-
         new_token_text = tokenizer.decode(next_token[0])
         output_text += new_token_text
-
         yield output_text
-
-        # Stop if EOS token generated
         if next_token.item() == tokenizer.eos_token_id:
             break
 
-def on_model_change(model_name):
-    status = load_model(model_name)
-    return status
+def respond(message, chat_history, reason_toggle, 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_label = "A1" if not use_fallback else "Fallback S2.1"
 
-with gr.Blocks() as demo:
-    gr.Markdown("# SmilyAI Sam Models — Manual Token Streaming Generator")
+    # Prefix prompt with reasoning mode
+    prefix = "/think " if reason_toggle else "/no_think "
+    processed_message = prefix + message.strip()
 
-    with gr.Row():
-        model_selector = gr.Dropdown(choices=MODELS, value=MODELS[0], label="Select Model")
-        status = gr.Textbox(label="Status", interactive=False)
+    chat_history = chat_history + [[message, ""]]
 
-    prompt_input = gr.Textbox(lines=3, placeholder="Enter your prompt here...", label="Prompt")
-    output_box = gr.Textbox(label="Generated Text", lines=15, interactive=False)
+    for response in generate_stream(processed_message, use_fallback=use_fallback):
+        chat_history[-1][1] = response + f" ({model_label})"
+        yield chat_history, chat_history
 
-    generate_btn = gr.Button("Generate")
+def clear_chat():
+    return [], []
 
-    # Load default model
-    status.value = load_model(MODELS[0])
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 SmilyAI Chatbot with Reasoning Toggle & Usage Limits")
+
+    model_status = gr.Textbox(label="Model Status", interactive=False)
+
+    chat_box = gr.Chatbot()
+    chat_history_state = gr.State([])
+
+    with gr.Row():
+        user_input = gr.Textbox(placeholder="Type your message...", show_label=False, scale=6)
+        reason_toggle = gr.Checkbox(label="Reason", value=True, scale=1)
+        send_btn = gr.Button("Send", scale=1)
 
-    model_selector.change(on_model_change, inputs=model_selector, outputs=status)
+    clear_btn = gr.Button("Clear Chat")
 
-    def generate_func(prompt):
-        if not prompt.strip():
-            yield "Please enter a prompt."
-            return
-        yield from generate_stream(prompt)
+    model_status.value = load_models()
 
-    generate_btn.click(generate_func, inputs=prompt_input, outputs=output_box)
+    send_btn.click(
+        respond,
+        inputs=[user_input, chat_history_state, reason_toggle],
+        outputs=[chat_box, chat_history_state]
+    )
 
-demo.launch()
+    clear_btn.click(fn=clear_chat, inputs=[], outputs=[chat_box, chat_history_state])