Update app.py

app.py

@@ -75,7 +75,7 @@ def model_prediction(model, text, device):
 # --- Llama 3.2 3B Model Setup ---
 LLAMA_MAX_MAX_NEW_TOKENS = 512  # Max tokens for Explanation
 LLAMA_DEFAULT_MAX_NEW_TOKENS = 512  # Max tokens for explantion
-LLAMA_MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "700")) # Reduced
+LLAMA_MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "2048")) # Reduced
 llama_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # Explicit device
 llama_model_id = "meta-llama/Llama-3.2-3B-Instruct"
 llama_tokenizer = AutoTokenizer.from_pretrained(llama_model_id)
@@ -125,28 +125,43 @@ def llama_generate(
 
 
 def generate_explanation(issue_text, top_qualities):
-    """Generates an explanation using Llama 3.2 3B."""
     if not top_qualities:
-        return "No explanation available as no quality tags were predicted."
+        return "<div style='color: red;'>No explanation available as no quality tags were predicted.</div>"
 
-    # Build the prompt, explicitly mentioning each quality
-    prompt_parts = [
-        "Given the following issue description:\n---\n",
-        issue_text,
-        "\n---\n",
-        "Explain why this issue might be classified under the following quality categories. Provide a concise explanation for each category, relating it back to the issue description:\n"
-    ]
-    for quality, _ in top_qualities:  # Iterate through qualities
-        prompt_parts.append(f"- {quality}\n")
+    prompt = f"""
+Given the following issue description:
+---
+{issue_text}
+---
+Explain why this issue might be classified under the following quality categories. Provide a concise explanation for each category, relating it back to the issue description:
+"""
+    for quality, _ in top_qualities:
+        prompt += f"- {quality}\n"
 
-    prompt = "".join(prompt_parts)
 
     try:
-        explanation = llama_generate(prompt)  # Get the explanation (not streamed)
-        return explanation
+        explanation = llama_generate(prompt)
+        # Format the explanation for better readability
+        formatted_explanation = ""
+        for quality, _ in top_qualities:
+            formatted_explanation += f"<p><b>{quality}:</b></p>"  # Bold the quality name
+            # Find the explanation for this specific quality.  This is a simple
+            # approach that works if Llama follows the prompt structure.
+            # A more robust approach might use regex or sentence embeddings.
+            start = explanation.find(quality)
+            if start != -1:
+                start += len(quality) + 2 # Move past "Quality:"
+                end = explanation.find("\n", start)  # Find next newline
+                if end == -1:
+                    end = len(explanation)
+                formatted_explanation += f"<p>{explanation[start:end].strip()}</p>" # Add the explanation text
+            else:
+                formatted_explanation += f"<p>Explanation for {quality} not found.</p>"
+
+        return f"<div style='overflow-y: scroll; max-height: 400px;'>{formatted_explanation}</div>" #Added scroll
     except Exception as e:
         logging.error(f"Error during Llama generation: {e}")
-        return "An error occurred while generating the explanation."
+        return "<div style='color: red;'>An error occurred while generating the explanation.</div>"
 
 
  # @spaces.GPU(duration=60)  # Apply the GPU decorator *only* to the main interface
@@ -229,10 +244,10 @@ interface = gr.Interface(
     outputs=[
         gr.HTML(label="Prediction Output"),
         gr.Textbox(label="Predictions", visible=False),
-        gr.Textbox(label="Explanation", lines=5)
+        gr.HTML(label="Explanation") # Change to gr.HTML
     ],
     title="QualityTagger",
     description="This tool classifies text into different quality domains such as Security, Usability, etc., and provides explanations.",
     examples=example_texts
 )
-interface.launch(share=True)
+interface.launch(share=True)