Merge pull request #9 from datalab-to/dev

Improve robustness

chandra/model/vllm.py

@@ -72,15 +72,16 @@ def generate_vllm(
                 temperature=temperature,
                 top_p=top_p,
             )
+            result = GenerationResult(
+                raw=completion.choices[0].message.content,
+                token_count=completion.usage.completion_tokens,
+                error=False,
+            )
         except Exception as e:
             print(f"Error during VLLM generation: {e}")
             return GenerationResult(raw="", token_count=0, error=True)
 
-        return GenerationResult(
-            raw=completion.choices[0].message.content,
-            token_count=completion.usage.completion_tokens,
-            error=False,
-        )
+        return result
 
     def process_item(item, max_retries):
         result = _generate(item)

chandra/settings.py

@@ -13,7 +13,7 @@ class Settings(BaseSettings):
     MIN_IMAGE_DIM: int = 1536
     MODEL_CHECKPOINT: str = "datalab-to/chandra"
     TORCH_DEVICE: str | None = None
-    MAX_OUTPUT_TOKENS: int = 8192
+    MAX_OUTPUT_TOKENS: int = 12384
     TORCH_ATTN: str | None = None
 
     # vLLM server settings