Update hf_model.py

hf_model.py

@@ -1,12 +1,19 @@
 from transformers import pipeline
 
 # Load the Hugging Face pipeline
-def load_model(task="summarization"):
-    model = pipeline(task=task, model="facebook/bart-large-cnn")
+def load_model(task="summarization", framework="pt"):
+    """
+    Load the specified task model using Hugging Face's pipeline.
+    Default is PyTorch ('pt') as the framework.
+    """
+    model = pipeline(task=task, model="facebook/bart-large-cnn", framework=framework)
     return model
 
 # Summarization function
 def summarize_text(model, text):
+    """
+    Summarize the provided legal text.
+    """
     if not text.strip():
         return "Please provide input text."
     result = model(text, max_length=150, min_length=40, do_sample=False)
@@ -14,5 +21,10 @@ def summarize_text(model, text):
 
 # Question Answering function
 def answer_question(model, question, context):
+    """
+    Answer a question based on the provided legal context.
+    """
+    if not context.strip() or not question.strip():
+        return "Please provide both a context and a question."
     result = model(question=question, context=context)
     return result['answer']