Update app.py

app.py

@@ -1,3 +1,8 @@
+"""
+Module for loading a LoRA fine-tuned BART model and serving
+an interactive Gradio interface for text generation.
+"""
+
 import torch
 import gradio as gr
 from transformers import AutoTokenizer
@@ -5,9 +10,14 @@ from transformers import BartForConditionalGeneration
 from peft import PeftModel
 
 
-def load_model():
+def load_model() -> tuple[AutoTokenizer, PeftModel, torch.device]:
     """
-    Load tokenizer, and the fine-tuned LoRA model.
+    Load tokenizer and LoRA-enhanced model onto available device.
+
+    Returns:
+        tokenizer (AutoTokenizer): Tokenizer for text processing.
+        model (PeftModel): Fine-tuned LoRA BART model in eval mode.
+        device (torch.device): Computation device (GPU if available, else CPU).
     """
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -18,7 +28,9 @@ def load_model():
     base_model.config.attn_implementation = "sdpa"
 
     # Load PEFT (LoRA) model for inference
-    model = PeftModel.from_pretrained(base_model, "outputs/bart-base-reddit-lora").eval()
+    model = PeftModel.from_pretrained(
+        base_model, "outputs/bart-base-reddit-lora"
+    ).eval()
     model.to(device)
     model.eval()
 
@@ -31,7 +43,13 @@ tokenizer, model, device = load_model()
 
 def predict(text: str) -> str:
     """
-    Generate a response for a single input text.
+    Generate a text response given an input prompt.
+
+    Args:
+        text (str): The input prompt string.
+
+    Returns:
+        str: The decoded model output.
     """
     # Tokenize and move inputs to device
     inputs = tokenizer(
@@ -44,8 +62,8 @@ def predict(text: str) -> str:
     # Generate with both beam search and sampling for diversity
     outputs = model.generate(
         **inputs,
-        max_length=500,
-        num_beams=5,
+        max_length=128,
+        num_beams=10,
         do_sample=True,
         length_penalty=1.2,
         repetition_penalty=1.3,
@@ -60,13 +78,16 @@ def predict(text: str) -> str:
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 
-def main():
+def main() -> None:
+    """
+    Launch Gradio web interface for interactive model inference.
+    """
     interface = gr.Interface(
         fn=predict,
-        inputs=gr.Textbox(lines=5, placeholder="Ask a Question", label="Your Question"),
-        outputs=gr.Textbox(label="Model Output"),
-        title="Bart-Reddit-LoRA Inference",
-        description="Enter your prompt and click Submit to get the model's response.",
+        inputs=gr.Textbox(lines=5, placeholder="Broad questions often have better results.", label="Your Question"),
+        outputs=gr.Textbox(label="Mimic Bot's Comment"),
+        title="Reddit-User-Mimic-Bot Inference (Bart-LoRA)",
+        description="Enter a question you would ask on reddit, and our Mimic Bot would comment back! Have fun.",
         allow_flagging="never",
     )
     interface.launch()