working without async

app.py

@@ -1,15 +1,15 @@
 import sys
 import warnings
 import spaces
-import asyncio
 from threading import Thread
-from transformers import AsyncTextIteratorStreamer
+from transformers import TextIteratorStreamer
 from functools import partial
 
 import gradio as gr
 import torch
 import numpy as np
 from model import Rank1
+import math
 
 print(f"NumPy version: {np.__version__}")
 print(f"PyTorch version: {torch.__version__}")
@@ -18,22 +18,102 @@ print(f"PyTorch version: {torch.__version__}")
 warnings.filterwarnings("ignore", category=UserWarning, message="Can't initialize NVML")
     
 @spaces.GPU
-async def process_input(query: str, passage: str, stream: bool = True) -> tuple[str, str, str]:
+def process_input(query: str, passage: str, stream: bool = True) -> tuple[str, str, str]:
     """Process input through the reranker and return formatted outputs."""
-    try:
-        reranker = Rank1(model_name_or_path="orionweller/rank1-32b-awq")
-        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"
+    reranker = Rank1(model_name_or_path="orionweller/rank1-32b-awq")
+    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>"
+
+    reranker.model = reranker.model.to("cuda")
+    inputs = reranker.tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=reranker.context_size
+    ).to("cuda")
+
+    if stream:
+        streamer = TextIteratorStreamer(
+            reranker.tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+        
+        current_text = "<think>"
+        generation_output = None
+        
+        def generate_with_output():
+            nonlocal generation_output
+            generation_output = reranker.model.generate(
+                **inputs,
+                generation_config=reranker.generation_config,
+                stopping_criteria=reranker.stopping_criteria,
+                return_dict_in_generate=True,
+                output_scores=True,
+                streamer=streamer
+            )
+        
+        thread = Thread(target=generate_with_output)
+        thread.start()
+        
+        # Stream tokens as they're generated
+        for new_text in streamer:
+            current_text += new_text
+            yield (
+                "Processing...",
+                "Processing...",
+                current_text
+            )
+        
+        thread.join()
+        
+        # Add the stopping sequence and calculate final scores
+        current_text += "\n" + reranker.stopping_criteria[0].matched_sequence
+        
+        with torch.no_grad():
+            final_scores = generation_output.scores[-1][0]
+            true_logit = final_scores[reranker.true_token].item()
+            false_logit = final_scores[reranker.false_token].item()
+            true_score = math.exp(true_logit)
+            false_score = math.exp(false_logit)
+            score = true_score / (true_score + false_score)
+        
+        yield (
+            score > 0.5,
+            score,
+            current_text
+        )
+    else:
+        # Non-streaming mode
+        with torch.no_grad():
+            outputs = reranker.model.generate(
+                **inputs,
+                generation_config=reranker.generation_config,
+                stopping_criteria=reranker.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[reranker.true_token].item()
+            false_logit = final_scores[reranker.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 = reranker.tokenizer.decode(new_text)
+            output_reasoning = "<think>\n" + decoded_input.strip() + f"\n</think> {'true' if score > 0.5 else 'false'}"
+            
+            yield (
+                "Relevant" if score > 0.5 else "Not Relevant",
+                f"{score:.2%}",
+                output_reasoning
+            )
 
 # Example inputs
 examples = [

model.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import logging
 import math
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, TextStreamer, AsyncTextIteratorStreamer, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, TextStreamer, TextIteratorStreamer
 from transformers import StoppingCriteria, StoppingCriteriaList
 from transformers import AwqConfig, AutoModelForCausalLM
 from threading import Thread
@@ -80,103 +80,3 @@ class Rank1:
             eos_token_id=self.tokenizer.eos_token_id,
             stopping_sequences=["</think> true", "</think> false"]
         )
-
-    def predict(self, query: str, passage: str, stream: bool = False):
-        """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 stream:
-            streamer = TextIteratorStreamer(
-                self.tokenizer,
-                skip_prompt=True,
-                skip_special_tokens=True
-            )
-            
-            current_text = "<think>"
-            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=streamer
-                )
-            
-            thread = Thread(target=generate_with_output)
-            thread.start()
-            
-            # Stream tokens as they're generated
-            for new_text in streamer:
-                current_text += new_text
-                yield {
-                    "is_relevant": None,
-                    "confidence_score": None,
-                    "model_reasoning": current_text
-                }
-            
-            thread.join()
-            
-            # Add the stopping sequence and calculate final scores
-            current_text += "\n" + self.stopping_criteria[0].matched_sequence
-            
-            with torch.no_grad():
-                final_scores = generation_output.scores[-1][0]
-                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()