Rewrite to work with score models

app.py

@@ -27,288 +27,86 @@ from huggingface_hub import login
 login(token=HF_LE_LLM_READ_TOKEN)
 
 # Constants
-# MODEL_ID = "le-llm/lapa-v0.1-reasoning-only-32768"
-# MODEL_ID = "le-llm/lapa-v0.1-instruct"
-# MODEL_ID = "le-llm/lapa-v0.1-matt-instruction-5e06"
-# MODEL_ID = "le-llm/lapa-v0.1-reprojected"
-# MODEL_ID = "le-llm/lapa-v0.1.1-instruct"
-MODEL_ID = "le-llm/manipulative-score-model"
+DEFAULT_MODEL = "le-llm/manipulative-score-model"
+DEVICE = "cuda"
 
-MAX_TOKENS = 4096
-TEMPERATURE = 0.7
-TOP_P = 0.95
 
-IMAGE_MAX_SIZE = 1024
+# --- Cache to avoid repeated reloads ---
+_model_cache: Dict[str, tuple[torch.nn.Module, AutoTokenizer]] = {}
 
 
-def _begin_analytics_session():
-    # Called once per client on app load
-    pass
-    # _ = logger.start_session(MODEL_ID)
+def load_model(model_id: str):
+    """Load model + tokenizer, auto-detect whether it's embedding or regression."""
+    if model_id in _model_cache:
+        return _model_cache[model_id]
 
-def load_model():
-    """Lazy-load model, tokenizer, and optional processor (for zeroGPU)."""
-    device = "cuda"  # if torch.cuda.is_available() else "cpu"
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    processor = None
-    try:
-        processor = AutoProcessor.from_pretrained(MODEL_ID)
-    except Exception as err:  # pragma: no cover - informative fallback
-        print(f"Warning: AutoProcessor not available ({err}). Falling back to tokenizer.")
+    print(f"🔹 Loading model: {model_id}")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-    model = AutoModel.from_pretrained(
-        MODEL_ID,
-        dtype=torch.bfloat16,  # if device == "cuda" else torch.float32,
-        device_map="cuda",  # if device == "cuda" else None,
-    )  # .cuda()
-    print(f"Selected device:", device)
-    return model, tokenizer, processor, device
+    model = AutoModel.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+    print(f"Detected embedding model: {model_id}")
 
+    model.to(DEVICE).eval()
+    _model_cache[model_id] = (model, tokenizer)
+    print(f"✅ Loaded model on {DEVICE}")
+    return model, tokenizer
 
-# Load model/tokenizer each request → allows zeroGPU to cold start & then release
-model, tokenizer, processor, device = load_model()
 
-
-def user(user_message, history: list):
-    """Format user message with optional image."""
-    import io
-
-    user_message = user_message or ""
-    updated_history = list(history)
-    has_content = False
-
-    stripped_message = user_message.strip()
-
-    if stripped_message:
-        has_content = True
-
-    if not has_content:
-        # Nothing to submit yet; keep inputs unchanged
-        return user_message, history
-
-    return "", updated_history
-
-
-def append_example_message(x: gr.SelectData, history):
-    if x.value["text"] is not None:
-        history.append({"role": "user", "content": x.value["text"]})
-
-    return history
-
-
-def _extract_text_from_content(content: Any) -> str | tuple[str, str]:
-    """Extract text from message content for logging."""
-    if isinstance(content, str):
-        return content
-    if isinstance(content, tuple) and len(content) == 2:
-        return content # (image_path, user_text)
-
-    raise ValueError(f"Unsupported content type for text extraction: {content}")
-
-
-def _clean_history_for_display(history: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    """Remove internal metadata fields like _base64 before displaying in Gradio."""
-    cleaned = []
-    for message in history:
-        cleaned_message = {"role": message.get("role", "user")}
-        content = message.get("content")
-
-        if isinstance(content, str):
-            cleaned_message["content"] = content
-        elif isinstance(content, list):
-            cleaned_content = []
-            for item in content:
-                if isinstance(item, dict):
-                    # Remove _base64 metadata
-                    cleaned_item = {k: v for k, v in item.items() if not k.startswith("_")}
-                    cleaned_content.append(cleaned_item)
-                else:
-                    cleaned_content.append(item)
-            cleaned_message["content"] = cleaned_content
-        else:
-            cleaned_message["content"] = content
-
-        cleaned.append(cleaned_message)
-
-    return cleaned
+# --- Helper for embeddings ---
+def average_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
 
 
+# --- Main scoring logic ---
 @spaces.GPU
-def bot(
-    input: list[dict[str, Any]]
-):
-    """Generate bot response with support for text."""
-
-    # Early return if no input
-    if not input:
-        return
-
-    clean_input = f"query: {input}"
-    batch_dict = tokenizer(clean_input, max_length=512, padding=True, truncation=True, return_tensors='pt')
-    batch_dict = {k: v.to(device) for k, v in batch_dict.items()}
-
-    outputs = model(**batch_dict)
-    embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
-
-    embeddings = F.normalize(embeddings, p=2, dim=1)
-    scores = (embeddings[:2] @ embeddings[2:].T) * 100
-    return str(scores.tolist())
-
+def bot(user_message: str, history: list[dict[str, Any]], model_choice: str):
+    if not user_message.strip():
+        return "", history
+
+    model, tokenizer = load_model(model_choice)  # returns embedding model
+    history = history + [{"role": "user", "content": user_message}]
+
+    batch = tokenizer([user_message], padding=True, truncation=True, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        outputs = model(**batch)
+        # outputs.last_hidden_state.shape = [batch_size, seq_len, hidden_dim]
+        # average pool over tokens
+        embedding = average_pool(outputs.last_hidden_state, batch["attention_mask"])
+        score = model.score_head(embedding).squeeze().item()
+        # embedding = F.normalize(embedding, p=2, dim=1)  # optional
+        #
+        # # Compute scalar score from embedding (example: mean of embedding dims)
+        # score = embedding.mean().item()
+
+    response = f"🔹 {model_choice} → score: {score:.4f}"
+    history.append({"role": "assistant", "content": response})
+    return "", history
+
+# --- UI ---
+THEME = gr.themes.Soft(primary_hue="blue", secondary_hue="amber", neutral_hue="stone")
+
+MODEL_OPTIONS = [
+    "le-llm/manipulative-score-model",
+    "le-llm/gec-score-model"
+]
 
-def average_pool(last_hidden_states: torch.Tensor,
-                 attention_mask: torch.Tensor) -> torch.Tensor:
-    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
-    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
-
-# --- drop-in UI compatible with older Gradio versions ---
-import os, tempfile, time
-import gradio as gr
+def _clear_chat():
+    return "", []
 
-# Ukrainian-inspired theme with deep, muted colors reflecting unbeatable spirit:
-THEME = gr.themes.Soft(
-    primary_hue="blue",      # Deep blue representing Ukrainian sky and resolve
-    secondary_hue="amber",   # Warm amber representing golden fields and determination  
-    neutral_hue="stone",     # Earthy stone representing strength and foundation
-)
 
-# Load CSS from external file
-def load_css():
-    try:
-        with open("static/style.css", "r", encoding="utf-8") as f:
-            return f.read()
-    except FileNotFoundError:
-        print("Warning: static/style.css not found")
-        return ""
+with gr.Blocks(theme=THEME, fill_height=True) as demo:
+    gr.Markdown("### 🤔 LAPA Quality Estimation")
 
-CSS = load_css()
+    with gr.Row():
+        model_choice = gr.Dropdown(MODEL_OPTIONS, value=DEFAULT_MODEL, label="Select Model")
 
-def _clear_chat():
-    return "", []
+    chatbot = gr.Chatbot(type="messages", height=480)
+    msg = gr.Textbox(label=None, placeholder="Type your text…", lines=1)
+    clear_btn = gr.Button("Clear")
 
-with gr.Blocks(theme=THEME, css=CSS, fill_height=True) as demo:
-    demo.load(fn=_begin_analytics_session, inputs=None, outputs=None)
-
-
-    # Header (no gr.Box to avoid version issues)
-    gr.HTML(
-        """
-        <div id="app-header">
-          <div class="app-title">🤔 LAPA Quality Estimation</div>
-        </div>
-        """
-    )
-
-    with gr.Row(equal_height=True):
-        # Left side: Chat
-        with gr.Column(scale=7, elem_id="left-pane"):
-            with gr.Column(elem_id="chat-card"):
-                chatbot = gr.Chatbot(
-                    type="messages",
-                    height=560,
-                    render_markdown=True,
-                    show_copy_button=True,
-                    show_label=False,
-                    # likeable=True,
-                    allow_tags=["think"],
-                    elem_id="chatbot",
-                    examples=[
-                        {"text": i}
-                        for i in [
-                            "хто тримає цей район?",
-                            "Напиши історію про Івасика-Телесика",
-                            "Яка найвища гора в Україні?",
-                            "Як звали батька Тараса Григоровича Шевченка?",
-                            "Яка з цих гір не знаходиться у Європі? Говерла, Монблан, Гран-Парадізо, Еверест",
-                            "Дай відповідь на питання\nЧому у качки жовті ноги?",
-                        ]
-                    ],
-                )
-
-
-            # ChatGPT-style input box with stop button
-            with gr.Row(elem_id="chat-input-row"):
-                msg = gr.Textbox(
-                    label=None,
-                    placeholder="Message… (Press Enter to send)",
-                    autofocus=True,
-                    lines=1,
-                    max_lines=6,
-                    container=False,
-                    show_label=False,
-                    elem_id="chat-input",
-                    elem_classes=["chat-input-box"]
-                )
-                stop_btn_visible = gr.Button(
-                    "⏹️", 
-                    variant="secondary", 
-                    elem_id="stop-btn-visible",
-                    elem_classes=["stop-btn-chat"],
-                    visible=False,
-                    size="sm"
-                )
-            
-            # Hidden buttons for functionality
-            with gr.Row(visible=True, elem_id="hidden-buttons"):
-                send_btn = gr.Button("Send", variant="primary", elem_id="send-btn")
-                stop_btn = gr.Button("Stop", variant="secondary", elem_id="stop-btn")
-                clear_btn = gr.Button("Clear", variant="secondary", elem_id="clear-btn")
-
-            # export_btn = gr.Button("Export chat (.md)", variant="secondary", elem_classes=["rounded-btn","secondary-btn"])
-            # exported_file = gr.File(label="", interactive=False, visible=True)
-            gr.HTML('<div class="footer-tip">Shortcuts: Enter to send • Shift+Enter for new line</div>')
-
-    # Helper functions for managing UI state
-    def show_stop_button():
-        return gr.update(visible=True)
-    
-    def hide_stop_button():
-        return gr.update(visible=False)
-
-    # Events (preserve your original handlers)
-    e1 = msg.submit(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=True).then(
-        fn=show_stop_button, inputs=None, outputs=stop_btn_visible
-    ).then(
-        fn=bot, inputs=chatbot, outputs=chatbot
-    ).then(
-        fn=hide_stop_button, inputs=None, outputs=stop_btn_visible
-    )
-
-    e2 = send_btn.click(fn=user, inputs=[msg, chatbot], outputs=[msg,chatbot], queue=True).then(
-        fn=show_stop_button, inputs=None, outputs=stop_btn_visible
-    ).then(
-        fn=bot, inputs=chatbot, outputs=chatbot
-    ).then(
-        fn=hide_stop_button, inputs=None, outputs=stop_btn_visible
-    )
-
-    e3 = chatbot.example_select(fn=append_example_message, inputs=[chatbot], outputs=[chatbot], queue=True).then(
-        fn=show_stop_button, inputs=None, outputs=stop_btn_visible
-    ).then(
-        fn=bot, inputs=chatbot, outputs=chatbot
-    ).then(
-        fn=hide_stop_button, inputs=None, outputs=stop_btn_visible
-    )
-
-    # Stop cancels running events (both buttons work)
-    stop_btn.click(fn=hide_stop_button, inputs=None, outputs=stop_btn_visible, cancels=[e1, e2, e3], queue=True)
-    stop_btn_visible.click(fn=hide_stop_button, inputs=None, outputs=stop_btn_visible, cancels=[e1, e2, e3], queue=True)
-
-    # Clear chat + input
-    clear_btn.click(fn=_clear_chat, inputs=None, outputs=[msg, chatbot])
-
-    # Export markdown
-    # export_btn.click(fn=_export_markdown, inputs=chatbot, outputs=exported_file)
-
-    # Load and inject external JavaScript
-    def load_javascript():
-        try:
-            with open("static/script.js", "r", encoding="utf-8") as f:
-                return f"<script>{f.read()}</script>"
-        except FileNotFoundError:
-            print("Warning: static/script.js not found")
-            return ""
-    
-    gr.HTML(load_javascript())
+    msg.submit(bot, inputs=[msg, chatbot, model_choice], outputs=[msg, chatbot])
+    clear_btn.click(_clear_chat, outputs=[msg, chatbot])
 
 if __name__ == "__main__":
     demo.queue().launch()