move away from async

model.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import logging
 import math
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, TextStreamer, AsyncTextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, TextStreamer, AsyncTextIteratorStreamer, TextIteratorStreamer
 from transformers import StoppingCriteria, StoppingCriteriaList
 from transformers import AwqConfig, AutoModelForCausalLM
 from threading import Thread
@@ -81,7 +81,7 @@ class Rank1:
             stopping_sequences=["</think> true", "</think> false"]
         )
 
-    async def predict(self, query: str, passage: str, streamer=None):
+    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" \
@@ -96,18 +96,16 @@ class Rank1:
             max_length=self.context_size
         ).to("cuda")
 
-        if streamer:
-            # Create a new streamer for each prediction
-            actual_streamer = AsyncTextIteratorStreamer(
+        if stream:
+            streamer = TextIteratorStreamer(
                 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(
@@ -116,14 +114,14 @@ class Rank1:
                     stopping_criteria=self.stopping_criteria,
                     return_dict_in_generate=True,
                     output_scores=True,
-                    streamer=actual_streamer
+                    streamer=streamer
                 )
             
             thread = Thread(target=generate_with_output)
             thread.start()
             
             # Stream tokens as they're generated
-            async for new_text in actual_streamer:
+            for new_text in streamer:
                 current_text += new_text
                 yield {
                     "is_relevant": None,
@@ -133,12 +131,11 @@ class Rank1:
             
             thread.join()
             
-            # Add the stopping sequence that was matched
+            # Add the stopping sequence and calculate final scores
             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
+                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)