Create app.py

app.py

@@ -0,0 +1,483 @@
+# -------------------------------
+# app.py (CPU COMPATIBLE VERSION)
+#
+# This file contains the backend logic and Gradio UI for the chatbot.
+#
+# --- FINAL, WORKING VERSION ---
+# - Specifies target_modules in LoraConfig to work with the custom Sam2 model.
+# - Uses a pure PyTorch fine-tuning loop for maximum control and stability.
+# - Custom Sam2Config inherits from PretrainedConfig to solve subscriptable errors.
+# - UI polling is backward-compatible with older Gradio versions.
+# -------------------------------
+import time
+import math
+import json
+import requests
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from pathlib import Path
+from safetensors.torch import load_file
+from transformers import AutoTokenizer, AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+import gradio as gr
+import os
+from datetime import datetime
+import threading
+import time
+import traceback
+
+# --- RLHF & Training Imports ---
+from huggingface_hub import HfApi, login
+from datasets import Dataset, load_dataset, concatenate_datasets
+from peft import LoraConfig, get_peft_model
+
+# -------------------------------
+# 0) RLHF & TUNING CONFIGURATION
+# -------------------------------
+FEEDBACK_DATASET_REPO = "Smilyai-labs/Open-Sam-2.5-chat"
+TUNED_MODEL_REPO_OWNER = "Smilyai-labs"
+BASE_MODEL_REPO = "Smilyai-labs/Sam-2.5-PRO-SOLVER-V2"
+FINETUNE_TRIGGER_LIKES = 2
+MIN_LIKES_FOR_TRAINING = 2
+
+# --- PyTorch Training Config ---
+LEARNING_RATE = 2e-4
+NUM_EPOCHS = 1
+BATCH_SIZE = 1
+
+# --- Login to Hugging Face Hub ---
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    print("WARNING: Hugging Face token not found. Feedback will not be saved and tuning will not run.")
+else:
+    login(token=HF_TOKEN)
+    print("Hugging Face token found. Feedback logging and model tuning are enabled.")
+
+# --- Global state ---
+LIKE_COUNTER = 0
+like_counter_lock = threading.Lock()
+training_lock = threading.Lock()
+model_lock = threading.Lock()
+TRAINING_STATUS = ""
+
+# -------------------------------
+# 1) Local Sam-2 architecture
+# -------------------------------
+class Sam2Config(PretrainedConfig):
+    model_type = "sam2"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        d_model=384,
+        n_layers=6,
+        n_heads=6,
+        ff_mult=4.0,
+        dropout=0.1,
+        input_modality="text",
+        head_type="causal_lm",
+        version="0.1",
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.ff_mult = ff_mult
+        self.dropout = dropout
+        self.input_modality = input_modality
+        self.head_type = head_type
+        self.version = version
+        super().__init__(**kwargs)
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d))
+    def forward(self, x):
+        return self.weight * x * (x.pow(2).mean(-1, keepdim=True) + self.eps).rsqrt()
+
+class MHA(nn.Module):
+    def __init__(self, d_model, n_heads, dropout=0.0):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        self.dropout = nn.Dropout(dropout)
+    def forward(self, x, attn_mask=None):
+        B, T, C = x.shape
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        causal = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
+        scores = scores.masked_fill(causal, float("-inf"))
+        if attn_mask is not None:
+            scores = scores.masked_fill(~attn_mask.unsqueeze(1).unsqueeze(2).bool(), float("-inf"))
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(self.dropout(attn), v).transpose(1, 2).contiguous().view(B, T, C)
+        return self.out_proj(out)
+
+class SwiGLU(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.w1 = nn.Linear(d_model, d_ff, bias=False)
+        self.w2 = nn.Linear(d_model, d_ff, bias=False)
+        self.w3 = nn.Linear(d_ff, d_model, bias=False)
+        self.dropout = nn.Dropout(dropout)
+    def forward(self, x):
+        return self.w3(self.dropout(torch.nn.functional.silu(self.w1(x)) * self.w2(x)))
+
+class Block(nn.Module):
+    def __init__(self, d_model, n_heads, ff_mult, dropout=0.0):
+        super().__init__()
+        self.norm1 = RMSNorm(d_model)
+        self.attn = MHA(d_model, n_heads, dropout=dropout)
+        self.norm2 = RMSNorm(d_model)
+        self.ff = SwiGLU(d_model, int(ff_mult * d_model), dropout=dropout)
+        self.drop = nn.Dropout(dropout)
+    def forward(self, x, attn_mask=None):
+        x = x + self.drop(self.attn(self.norm1(x), attn_mask=attn_mask))
+        x = x + self.drop(self.ff(self.norm2(x)))
+        return x
+
+class Sam2(PreTrainedModel):  # <-- CHANGE THIS LINE: inherit from PreTrainedModel
+    config_class = Sam2Config  # <-- ADD THIS LINE: tell HF what config class to use
+
+    def __init__(self, config: Sam2Config):
+        super().__init__(config)  # <-- CHANGE THIS LINE: pass config to parent
+        self.config = config  # You can keep this if you use it elsewhere
+        self.embed = nn.Embedding(config.vocab_size, config.d_model)
+        self.blocks = nn.ModuleList([Block(config.d_model, config.n_heads, config.ff_mult, dropout=config.dropout) for _ in range(config.n_layers)])
+        self.norm = RMSNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.embed.weight
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    def forward(self, input_ids=None, inputs_embeds=None, attention_mask=None, labels=None, **kwargs):
+        if inputs_embeds is not None:
+            x = inputs_embeds
+        else:
+            if input_ids is None:
+                raise ValueError("You must provide either input_ids or inputs_embeds")
+            x = self.embed(input_ids)
+
+        for blk in self.blocks:
+            x = blk(x, attn_mask=attention_mask)
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+        if loss is not None:
+            return (loss, logits)
+        return (logits,)
+
+# -------------------------------
+# 2) Load initial resources
+# -------------------------------
+weights_filename = "model.safetensors"
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_REPO)
+tokenizer.pad_token = tokenizer.eos_token
+# --- FIXED: Removed extra spaces in URLs ---
+config_url = f"https://huggingface.co/{BASE_MODEL_REPO}/raw/main/config.json"
+config_data = requests.get(config_url).json()
+cfg = Sam2Config(**config_data)
+
+# --- FIXED: Removed extra spaces in URLs ---
+weights_url = f"https://huggingface.co/{BASE_MODEL_REPO}/resolve/main/{weights_filename}"
+weights_content = requests.get(weights_url).content
+with open(weights_filename, "wb") as f: f.write(weights_content)
+
+model = Sam2(cfg)
+state_dict = load_file(weights_filename)
+model.load_state_dict(state_dict)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device).eval()
+print(f"Inference will run on: {device}")
+
+EOT_ID = tokenizer.convert_tokens_to_ids("<|eot|>") or tokenizer.eos_token_id
+SPECIAL_TOKENS = {"bos": "<|bos|>", "eot": "<|eot|>", "user": "<|user|>", "assistant": "<|assistant|>", "system": "<|system|>"}
+SYSTEM_PROMPT = "You are Sam-2, a friendly and concise chatbot. Always give short, direct answers and avoid medical or legal advice."
+
+AutoModelForCausalLM.register(Sam2Config, Sam2)
+
+# -------------------------------
+# 3) Inference and Feedback Functions
+# -------------------------------
+def sample_next_token( logits, past_tokens, temperature=0.8, top_k=40, top_p=0.9, repetition_penalty=1.1, max_repeat=5, no_repeat_ngram_size=3 ):
+    if logits.dim() == 3: logits = logits[:, -1, :].clone()
+    else: logits = logits.clone()
+    batch_size, vocab_size = logits.size(0), logits.size(1)
+    orig_logits = logits.clone()
+    if temperature != 1.0: logits = logits / float(temperature)
+    past_list = past_tokens.tolist() if isinstance(past_tokens, torch.Tensor) else list(past_tokens)
+    for token_id in set(past_list):
+        if 0 <= token_id < vocab_size: logits[:, token_id] /= repetition_penalty
+    if len(past_list) >= max_repeat:
+        last_token, count = past_list[-1], 1
+        for i in reversed(past_list[:-1]):
+            if i == last_token: count += 1
+            else: break
+        if count >= max_repeat: logits[:, last_token] = -float("inf")
+    if no_repeat_ngram_size > 0 and len(past_list) >= no_repeat_ngram_size:
+        ngram = tuple(past_list[-no_repeat_ngram_size:])
+        for token_id in range(vocab_size):
+            if tuple(past_list[-(no_repeat_ngram_size - 1):] + [token_id]) == ngram: logits[:, token_id] = -float("inf")
+    if top_k is not None and top_k > 0:
+        tk = min(max(1, int(top_k)), vocab_size)
+        topk_vals, _ = torch.topk(logits, tk, dim=-1)
+        min_topk = topk_vals[:, -1].unsqueeze(-1)
+        logits[logits < min_topk] = -float("inf")
+    if top_p is not None and 0.0 < top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        sorted_probs = F.softmax(sorted_logits, dim=-1)
+        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+        for b in range(batch_size):
+            sorted_mask = cumulative_probs[b] > top_p
+            if sorted_mask.numel() > 0:
+                sorted_mask[0] = False
+                tokens_to_remove = sorted_indices[b][sorted_mask]
+                logits[b, tokens_to_remove] = -float("inf")
+    for b in range(batch_size):
+        if torch.isneginf(logits[b]).all(): logits[b] = orig_logits[b]
+    probs = F.softmax(logits, dim=-1)
+    if torch.isnan(probs).any(): probs = torch.ones_like(logits) / logits.size(1)
+    next_token = torch.multinomial(probs, num_samples=1)
+    return next_token.to(device)
+
+def predict(message, history):
+    chat_history = []
+    for human, assistant in history:
+        chat_history.append(f"{SPECIAL_TOKENS['user']} {human} {SPECIAL_TOKENS['eot']}")
+        if assistant:
+            chat_history.append(f"{SPECIAL_TOKENS['assistant']} {assistant} {SPECIAL_TOKENS['eot']}")
+    chat_history.append(f"{SPECIAL_TOKENS['user']} {message} {SPECIAL_TOKENS['eot']}")
+    prompt = f"{SPECIAL_TOKENS['system']} {SYSTEM_PROMPT} {SPECIAL_TOKENS['eot']}\n" + "\n".join(chat_history) + f"\n{SPECIAL_TOKENS['assistant']}"
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+    generated_text = ""
+    for _ in range(256):
+        with torch.no_grad(), model_lock:
+            outputs = model(input_ids, attention_mask=attention_mask)
+        logits = outputs[0]
+        next_token = sample_next_token(logits, input_ids[0], temperature=0.4, top_k=50, top_p=0.9, repetition_penalty=1.1)
+        token_id = int(next_token.squeeze().item())
+        if token_id == EOT_ID: break
+        token_str = tokenizer.decode([token_id], skip_special_tokens=True)
+        input_ids = torch.cat([input_ids, next_token], dim=1)
+        attention_mask = torch.cat([attention_mask, torch.ones((attention_mask.size(0), 1), device=device, dtype=attention_mask.dtype)], dim=1)
+        generated_text += token_str
+        yield generated_text
+
+def log_feedback(data: gr.LikeData, history: list):
+    global LIKE_COUNTER
+    if not HF_TOKEN:
+        print("Feedback not logged. HF_TOKEN not set.")
+        return
+    feedback_entry = { "prompt": history[data.index[0]][0], "response": data.value, "feedback": 1 if data.liked else 0, "timestamp": datetime.utcnow().isoformat() }
+    new_feedback_dataset = Dataset.from_dict({k: [v] for k, v in feedback_entry.items()})
+    try:
+        existing_dataset = load_dataset(FEEDBACK_DATASET_REPO, split="train", cache_dir="./cache")
+        combined_dataset = concatenate_datasets([existing_dataset, new_feedback_dataset])
+    except Exception as e:
+        print(f"Could not load existing dataset: {e}. Creating a new one.")
+        combined_dataset = new_feedback_dataset
+    try:
+        combined_dataset.push_to_hub(FEEDBACK_DATASET_REPO, private=False)
+        feedback_icon = '👍' if data.liked else '👎'
+        print(f"Successfully logged {feedback_icon} feedback. Dataset now has {len(combined_dataset)} entries.")
+        if data.liked:
+            with like_counter_lock:
+                LIKE_COUNTER += 1
+                current_likes = LIKE_COUNTER
+            print(f"Like recorded. Total likes since start: {current_likes}.")
+            if current_likes > 0 and current_likes % FINETUNE_TRIGGER_LIKES == 0:
+                print(f"--- Like threshold of {FINETUNE_TRIGGER_LIKES} reached! Triggering fine-tuning. ---")
+                tuning_thread = threading.Thread(target=run_tuning_task, daemon=True)
+                tuning_thread.start()
+    except Exception as e:
+        print(f"Error logging feedback to Hub: {e}")
+
+
+# -------------------------------
+# 6) Background Fine-Tuning Logic (PyTorch Loop)
+# -------------------------------
+def run_tuning_task():
+    global model, TRAINING_STATUS
+
+    if not training_lock.acquire(blocking=False):
+        print("Tuning is already in progress. Skipping this trigger.")
+        return
+
+    print("\n--- Starting PyTorch Fine-Tuning Task ---")
+    try:
+        TRAINING_STATUS = "🔧 Preparing to improve Sam-2.5..."
+        
+        if not HF_TOKEN:
+            TRAINING_STATUS = "Error: HF_TOKEN not set. Cannot run tuning."
+            time.sleep(10)
+            return
+
+        feedback_data = load_dataset(FEEDBACK_DATASET_REPO, split="train", cache_dir="./cache")
+        liked_data = feedback_data.filter(lambda x: x['feedback'] == 1)
+        print(f"Found {len(liked_data)} total liked responses for training.")
+
+        if len(liked_data) < MIN_LIKES_FOR_TRAINING:
+            TRAINING_STATUS = f"✅ Improvement complete! (Not enough new data to train, will try again later)."
+            time.sleep(5)
+            return
+        
+        def format_for_training(example):
+            return { "text": f"{SPECIAL_TOKENS['system']} {SYSTEM_PROMPT} {SPECIAL_TOKENS['eot']}\n{SPECIAL_TOKENS['user']} {example['prompt']} {SPECIAL_TOKENS['eot']}\n{SPECIAL_TOKENS['assistant']} {example['response']} {SPECIAL_TOKENS['eot']}"}
+        train_dataset = liked_data.map(format_for_training)
+
+        print("Loading base model for tuning...")
+        model_to_tune = Sam2(cfg)
+        state_dict_to_tune = load_file(weights_filename)
+        model_to_tune.load_state_dict(state_dict_to_tune)
+        
+        # --- THIS IS THE FIX ---
+        # We explicitly tell PEFT which linear layers in our MHA block to adapt.
+        peft_config = LoraConfig(
+            r=16, 
+            lora_alpha=32, 
+            lora_dropout=0.05, 
+            bias="none", 
+            task_type="CAUSAL_LM",
+            target_modules=["q_proj", "v_proj"]
+        )
+        # --- END FIX ---
+        
+        peft_model = get_peft_model(model_to_tune, peft_config)
+        peft_model.to(device)
+        peft_model.print_trainable_parameters()
+
+        tokenized_dataset = train_dataset.map(lambda examples: tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512), batched=True)
+        # --- ADDED: Remove the unused 'text' column to clean up the dataset ---
+        tokenized_dataset = tokenized_dataset.remove_columns(["text"])
+        tokenized_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask'])
+        train_dataloader = DataLoader(tokenized_dataset, batch_size=BATCH_SIZE)
+
+        optimizer = torch.optim.AdamW(peft_model.parameters(), lr=LEARNING_RATE)
+
+        TRAINING_STATUS = f"🔧 Sam-2.5 is starting training on {len(liked_data)} examples... Thank you all for your contribution to the dataset. The model will train and hot swap shortly.(This can be slow on CPU)"
+        print("Starting model tuning on CPU...")
+        peft_model.train()
+        for epoch in range(NUM_EPOCHS):
+            time.sleep(0.01)
+            for i, batch in enumerate(train_dataloader):
+                input_ids = batch['input_ids'].to(device)
+                attention_mask = batch['attention_mask'].to(device)
+                outputs = peft_model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
+                loss = outputs[0]
+                loss.backward()
+                optimizer.step()
+                optimizer.zero_grad()
+                current_loss = loss.item()
+                print(f"Epoch {epoch+1}, Batch {i+1}/{len(train_dataloader)}, Loss: {current_loss:.4f}")
+                # --- UPDATE UI WITH LIVE LOSS ---
+                TRAINING_STATUS = f"🔧  You are witnessing the training of sam2.5. Training... Batch {i+1}/{len(train_dataloader)}, Loss: {current_loss:.4f}"
+
+        print("Tuning complete.")
+        
+        TRAINING_STATUS = "✨ Finishing up... Merging improvements."
+        merged_model = peft_model.merge_and_unload()
+        
+        # --- FIXED: Safe Model Swap using model_lock ---
+        with model_lock:
+            print("Hot-swapping live model...")
+            # Create a new instance and copy state, preserving the object reference
+            new_state_dict = merged_model.state_dict()
+            model.load_state_dict(new_state_dict)
+            model.to(device).eval()
+        
+        date_str = datetime.now().strftime("%Y%m%d-%H%M")
+        new_repo_id = f"{TUNED_MODEL_REPO_OWNER}/Sam-2.5-PUBLIC-RLHF-{date_str}"
+
+        print(f"Saving and uploading tuned model to {new_repo_id}...")
+
+        # Create a directory to save the model
+        local_dir = f"./{new_repo_id.split('/')[-1]}"
+        os.makedirs(local_dir, exist_ok=True)
+
+        # Save model using Hugging Face format
+        merged_model.save_pretrained(local_dir, safe_serialization=False)
+        tokenizer.save_pretrained(local_dir)
+
+        # Push to Hub
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.create_repo(repo_id=new_repo_id, repo_type="model", exist_ok=True)
+        api.upload_folder(
+            folder_path=local_dir,
+            repo_id=new_repo_id,
+            repo_type="model"
+        )
+
+        # Clean up local files
+        import shutil
+        shutil.rmtree(local_dir)
+
+        print("Upload and hot-swap complete!")
+        TRAINING_STATUS = "✅ Sam-2.5 has been successfully upgraded! Thank you. You have helped shaped the newest generation of sam 2.5 pro solver. You, helped make AI"
+        time.sleep(5)
+
+    except Exception as e:
+        print(f"An error occurred during the tuning process: {e}")
+        traceback.print_exc()
+        TRAINING_STATUS = f"An error occurred during training: {e}"
+        time.sleep(10)
+    finally:
+        TRAINING_STATUS = ""
+        training_lock.release()
+        print("--- PyTorch Fine-Tuning Task Finished ---")
+
+# -------------------------------
+# 7) UI Functions & Gradio Interface
+# -------------------------------
+def check_training_status():
+    global TRAINING_STATUS
+    if TRAINING_STATUS:
+        return gr.update(value=TRAINING_STATUS, visible=True)
+    else:
+        return gr.update(value="", visible=False)
+
+def poll_status_updater():
+    while True:
+        yield check_training_status()
+        time.sleep(1)
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="blue")) as demo:
+    gr.Markdown("""
+        # Sam-2.5-PRO-SOLVER-V2 Chat
+        A self-improving chatbot powered by Sam-2. Use the thumb icons to rate responses!
+        The model automatically fine-tunes on your positive feedback and gets smarter live.
+        """)
+    
+    training_status_md = gr.Markdown(value="", visible=False)
+    chatbot = gr.Chatbot(label="Sam-2", bubble_full_width=False)
+    chat_interface = gr.ChatInterface(fn=predict, chatbot=chatbot)
+    chatbot.like(log_feedback, inputs=[chatbot], outputs=None)
+    
+    demo.load(poll_status_updater, None, training_status_md)
+TRAINING_STATUS = "Not training yet. Waiting for more examples. Sam-2.5 is ready."
+if __name__ == "__main__":
+    print("Starting Gradio app. Tuning will be triggered by user feedback.")
+    demo.launch(show_api=True)