Update README.md

README.md

@@ -38,7 +38,7 @@ def generate_prompts_for_classification(article: str, summary_sentences: List[st
     return prompts
 
 def predict_with_hf_classification_pipeline(prompts: List[Dict], model_name: str, max_context_length: int = 512,
-                                            batch_size: int = 2):
+                                            batch_size: int = 2) -> List[str]:
     device = "cuda" if torch.cuda.is_available() else "cpu"
     text_classification_pipeline = pipeline("text-classification", model=model_name, device=device,
                                             batch_size=batch_size)