merged

README.md

@@ -1,12 +1,13 @@
 ---
-title: Rank1 Demo
-emoji: 🐠
-colorFrom: red
-colorTo: pink
+title: Rank1: Test Time Compute in Reranking
+emoji: 🏆
+colorFrom: blue
+colorTo: green
 sdk: gradio
 sdk_version: 5.17.1
 app_file: app.py
 pinned: false
+license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,105 @@
+import sys
+import warnings
+import asyncio
+from threading import Thread
+from transformers import AsyncTextIteratorStreamer
+from functools import partial
+
+import gradio as gr
+import torch
+import numpy as np
+from model import Rank1
+
+print(f"NumPy version: {np.__version__}")
+print(f"PyTorch version: {torch.__version__}")
+
+# Suppress CUDA initialization warning
+warnings.filterwarnings("ignore", category=UserWarning, message="Can't initialize NVML")
+
+try:
+    reranker = Rank1(model_name_or_path="orionweller/rank1-32b-awq")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    sys.exit(1)
+    
+@spaces.GPU
+async def process_input(query: str, passage: str, stream: bool = True) -> tuple[str, str, str]:
+    """Process input through the reranker and return formatted outputs."""
+    try:
+        async for result in reranker.predict(query, passage, streamer=stream):
+            if result["is_relevant"] is None:
+                # Intermediate streaming result
+                yield "Processing...", "Processing...", result["model_reasoning"]
+            else:
+                # Final result
+                relevance = "Relevant" if result["is_relevant"] else "Not Relevant"
+                confidence = f"{result['confidence_score']:.2%}"
+                reasoning = result["model_reasoning"]
+                yield relevance, confidence, reasoning
+    except Exception as e:
+        yield f"Error: {str(e)}", "N/A", "An error occurred during processing"
+
+# Example inputs
+examples = [
+    [
+        "What movies were directed by James Cameron?",
+        "Avatar: The Way of Water is a 2022 American epic science fiction film directed by James Cameron.",
+    ],
+    [
+        "What are the symptoms of COVID-19?",
+        "Common symptoms of COVID-19 include fever, cough, fatigue, loss of taste or smell, and difficulty breathing.",
+    ]
+]
+
+theme = gr.themes.Soft(
+    primary_hue="indigo",
+    font=["Inter", "ui-sans-serif", "system-ui", "sans-serif"],
+    neutral_hue="slate",
+    radius_size="lg",
+)
+
+with gr.Blocks(theme=theme, css=".red-text { color: red; }") as demo:
+    gr.Markdown("# Rank1: Test Time Compute in Reranking")
+    gr.HTML('NOTE: for demo purposes this is a <span style="color: red;">quantized</span> model with a <span style="color: red;">1024</span> context length. HF spaces cannot use vLLM so this is <span style="color: red;">significantly slower</span>')
+    
+    with gr.Row():
+        with gr.Column():
+            query_input = gr.Textbox(
+                label="Query",
+                placeholder="Enter your search query here",
+                lines=2
+            )
+            passage_input = gr.Textbox(
+                label="Passage",
+                placeholder="Enter the passage to check for relevance",
+                lines=6
+            )
+            submit_button = gr.Button("Check Relevance")
+        
+        with gr.Column():
+            relevance_output = gr.Textbox(label="Relevance")
+            confidence_output = gr.Textbox(label="Confidence")
+            reasoning_output = gr.Textbox(
+                label="Model Reasoning",
+                lines=10,
+                interactive=False
+            )
+    
+    gr.Examples(
+        examples=examples,
+        inputs=[query_input, passage_input],
+        outputs=[relevance_output, confidence_output, reasoning_output],
+        fn=partial(process_input, stream=False),
+        cache_examples=True,
+    )
+    
+    submit_button.click(
+        fn=process_input,
+        inputs=[query_input, passage_input],
+        outputs=[relevance_output, confidence_output, reasoning_output],
+        api_name="predict",
+        queue=True
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

model.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import logging
+import math
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, TextStreamer, AsyncTextIteratorStreamer
+from transformers import StoppingCriteria, StoppingCriteriaList
+from transformers import AwqConfig, AutoModelForCausalLM
+from threading import Thread
+
+
+logger = logging.getLogger(__name__)
+
+
+class ThinkStoppingCriteria(StoppingCriteria):
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.true_sequence = tokenizer("</think> true").input_ids[1:]  # Skip first token
+        self.false_sequence = tokenizer("</think> false").input_ids[1:]  # Skip first token
+        self.matched_sequence = None
+        
+    def __call__(self, input_ids, scores, **kwargs):
+        for sequence in [self.true_sequence, self.false_sequence]:
+            if input_ids.shape[1] >= len(sequence):
+                if all((input_ids[0, -(len(sequence)-i)] == sequence[i] for i in range(len(sequence)))):
+                    self.matched_sequence = "</think> true" if sequence is self.true_sequence else "</think> false"
+                    return True
+        return False
+
+
+class Rank1:
+    def __init__(
+        self,
+        model_name_or_path: str = "",
+        # set these just for demo, typically longer
+        context_size: int = 4000,
+        max_output_tokens: int = 1024,
+        **kwargs,
+    ):
+        self.context_size = context_size
+        self.max_output_tokens = max_output_tokens
+
+        # Initialize tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+        self.tokenizer.padding_side = "left"
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        # Cache commonly used token IDs
+        self.true_token = self.tokenizer(" true", add_special_tokens=False).input_ids[0]
+        self.false_token = self.tokenizer(" false", add_special_tokens=False).input_ids[0]
+
+        # Load AWQ model on CPU initially
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name_or_path,
+            device_map="cpu",
+            trust_remote_code=True,
+            attn_implementation="flash_attention_2" 
+        )
+
+        self.stopping_criteria = StoppingCriteriaList([
+            ThinkStoppingCriteria(self.tokenizer)
+        ])
+        
+        # Update generation config
+        self.generation_config = GenerationConfig(
+            max_new_tokens=max_output_tokens,
+            do_sample=False,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id
+        )
+
+        # Create text streamer
+        self.streamer = TextStreamer(self.tokenizer)
+
+        # Simple generation config
+        self.generation_config = GenerationConfig(
+            max_new_tokens=max_output_tokens,
+            do_sample=False,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            stopping_sequences=["</think> true", "</think> false"]
+        )
+
+    async def predict(self, query: str, passage: str, streamer=None):
+        """Predict relevance of passage to query."""
+        prompt = f"Determine if the following passage is relevant to the query. Answer only with 'true' or 'false'.\n" \
+                f"Query: {query}\n" \
+                f"Passage: {passage}\n" \
+                "<think>"
+
+        self.model = self.model.to("cuda")
+        inputs = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=self.context_size
+        ).to("cuda")
+
+        if streamer:
+            # Create a new streamer for each prediction
+            actual_streamer = AsyncTextIteratorStreamer(
+                self.tokenizer,
+                skip_prompt=True,
+                skip_special_tokens=True
+            )
+            
+            current_text = "<think>"
+            
+            # Run generation in a separate thread and store the output
+            generation_output = None
+            def generate_with_output():
+                nonlocal generation_output
+                generation_output = self.model.generate(
+                    **inputs,
+                    generation_config=self.generation_config,
+                    stopping_criteria=self.stopping_criteria,
+                    return_dict_in_generate=True,
+                    output_scores=True,
+                    streamer=actual_streamer
+                )
+            
+            thread = Thread(target=generate_with_output)
+            thread.start()
+            
+            # Stream tokens as they're generated
+            async for new_text in actual_streamer:
+                current_text += new_text
+                yield {
+                    "is_relevant": None,
+                    "confidence_score": None,
+                    "model_reasoning": current_text
+                }
+            
+            thread.join()
+            
+            # Add the stopping sequence that was matched
+            current_text += "\n" + self.stopping_criteria[0].matched_sequence
+            
+            # Calculate final scores using the last scores from generation
+            with torch.no_grad():
+                final_scores = generation_output.scores[-1][0]  # Get logits from last position
+                true_logit = final_scores[self.true_token].item()
+                false_logit = final_scores[self.false_token].item()
+                true_score = math.exp(true_logit)
+                false_score = math.exp(false_logit)
+                score = true_score / (true_score + false_score)
+            
+            yield {
+                "is_relevant": score > 0.5,
+                "confidence_score": score,
+                "model_reasoning": current_text
+            }
+        else:
+            # Non-streaming mode
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    generation_config=self.generation_config,
+                    stopping_criteria=self.stopping_criteria,
+                    return_dict_in_generate=True,
+                    output_scores=True
+                )
+                
+                # Get final score from generation outputs
+                final_scores = outputs.scores[-1][0]  # Get logits from last position
+                true_logit = final_scores[self.true_token].item()
+                false_logit = final_scores[self.false_token].item()
+                true_score = math.exp(true_logit)
+                false_score = math.exp(false_logit)
+                score = true_score / (true_score + false_score)
+
+                # only decode the generated text
+                new_text = outputs.sequences[0][len(inputs.input_ids[0]):]
+                decoded_input = self.tokenizer.decode(new_text)
+                output_reasoning = "<think>\n" + decoded_input.strip() + f"\n</think> {'true' if score > 0.5 else 'false'}"
+                
+                yield {
+                    "is_relevant": score > 0.5,
+                    "confidence_score": score,
+                    "model_reasoning": output_reasoning
+                }
+
+        # Move model back to CPU
+        self.model = self.model.to("cpu")
+        torch.cuda.empty_cache()

requirements.txt

@@ -0,0 +1,8 @@
+gradio==5.17.1
+spaces
+transformers==4.49.0
+numpy==1.24.3
+flash_attn==2.6.3
+autoawq==0.2.1
+autoawq_kernels==0.0.9
+torch==2.5.1+cu121