refactor: remove FLASHINFER environment variable and update LLM initialization for batch processing

generate_summaries_uv.py

@@ -20,7 +20,6 @@ from typing import Optional
 
 # Set environment variables to speed up model loading
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-# os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER"
 
 import polars as pl
 from datasets import Dataset, load_dataset
@@ -112,7 +111,10 @@ def generate_summaries(
 
     # Initialize model and tokenizer from local path
     logger.info(f"Initializing vLLM model from local path: {local_model_path}")
-    llm = LLM(model=local_model_path)
+    llm = LLM(
+        model=local_model_path,
+        max_model_len=4096,  # Adjust based on model capabilities
+    )
     tokenizer = AutoTokenizer.from_pretrained(local_model_path)
     sampling_params = SamplingParams(
         temperature=temperature,
@@ -131,10 +133,22 @@ def generate_summaries(
     logger.info(f"Generating summaries for {len(prompts)} items")
     all_outputs = []
 
-    for i in tqdm(range(0, len(prompts), batch_size), desc="Generating summaries"):
-        batch_prompts = prompts[i : i + batch_size]
-        outputs = llm.generate(batch_prompts, sampling_params)
-        all_outputs.extend(outputs)
+    # for i in tqdm(range(0, len(prompts), batch_size), desc="Generating summaries"):
+    #     batch_prompts = prompts[i : i + batch_size]
+    #     outputs = llm.generate(batch_prompts, sampling_params)
+    #     all_outputs.extend(outputs)
+    # try directly doing whole dataset
+    all_outputs = llm.generate(
+        prompts,
+        sampling_params,
+        batch_size=batch_size,
+        max_batch_size=batch_size,
+    )
+    logger.info(f"Generated {len(all_outputs)} summaries")
+    if len(all_outputs) != len(prompts):
+        logger.warning(
+            f"Generated {len(all_outputs)} summaries, but expected {len(prompts)}. Some prompts may have failed."
+        )
 
     # Extract clean results
     clean_results = [output.outputs[0].text.strip() for output in all_outputs]