Update app.py

app.py

@@ -53,14 +53,14 @@ quality_mapping = {
 
 # Pre-load models and tokenizer for quality prediction
 tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
-models = {path: AutoModelForSequenceClassification.from_pretrained(path) for path in model_paths}  # Load to CPU initially
+models = {path: AutoModelForSequenceClassification.from_pretrained(path) for path in model_paths}
 
 def get_quality_name(model_name):
     return quality_mapping.get(model_name.split('/')[-1], "Unknown Quality")
 
 
 def model_prediction(model, text, device):
-    model.to(device) # Move the *specific* model to the GPU
+    model.to(device)
     model.eval()
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     inputs = {k: v.to(device) for k, v in inputs.items()}
@@ -69,19 +69,19 @@ def model_prediction(model, text, device):
         logits = outputs.logits
         probs = softmax(logits.cpu().numpy(), axis=1)
     avg_prob = np.mean(probs[:, 1])
-    model.to("cpu")  # Move the model *back* to the CPU
+    model.to("cpu")
     return avg_prob
 
 # --- 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", "2048")) # Reduced
-llama_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # Explicit device
+LLAMA_MAX_MAX_NEW_TOKENS = 512
+LLAMA_DEFAULT_MAX_NEW_TOKENS = 512
+LLAMA_MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "1024"))
+llama_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 llama_model_id = "meta-llama/Llama-3.2-3B-Instruct"
 llama_tokenizer = AutoTokenizer.from_pretrained(llama_model_id)
 llama_model = AutoModelForCausalLM.from_pretrained(
     llama_model_id,
-    device_map="auto",  # Let Transformers handle optimal device placement
+    device_map="auto",
     torch_dtype=torch.bfloat16,
 )
 llama_model.eval()
@@ -96,7 +96,7 @@ def llama_generate(
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
-) -> str: # Return string, not iterator
+) -> str:
 
     inputs = llama_tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=LLAMA_MAX_INPUT_TOKEN_LENGTH).to(llama_model.device)
 
@@ -104,8 +104,7 @@ def llama_generate(
         inputs.input_ids = inputs.input_ids[:, -LLAMA_MAX_INPUT_TOKEN_LENGTH:]
         gr.Warning(f"Trimmed input from conversation as it was longer than {LLAMA_MAX_INPUT_TOKEN_LENGTH} tokens.")
 
-    # Generate *without* streaming
-    with torch.no_grad():  # Ensure no gradient calculation
+    with torch.no_grad():
         generate_ids = llama_model.generate(
             **inputs,
             max_new_tokens=max_new_tokens,
@@ -115,56 +114,41 @@ def llama_generate(
             temperature=temperature,
             num_beams=1,
             repetition_penalty=repetition_penalty,
-            pad_token_id=llama_tokenizer.pad_token_id,  # Pass pad_token_id here
-            eos_token_id=llama_tokenizer.eos_token_id,  # Pass eos_token_id here
+            pad_token_id=llama_tokenizer.pad_token_id,
+            eos_token_id=llama_tokenizer.eos_token_id,
 
         )
     output_text = llama_tokenizer.decode(generate_ids[0], skip_special_tokens=True)
-    torch.cuda.empty_cache()  # Clear cache after each generation
+    torch.cuda.empty_cache()
     return output_text
 
 
-def generate_explanation(issue_text, top_qualities):
-    if not top_qualities:
-        return "<div style='color: red;'>No explanation available as no quality tags were predicted.</div>"
+def generate_explanation(issue_text, top_quality):
+    """Generates an explanation for the *single* top quality above threshold."""
+    if not top_quality:
+        return "<div style='color: red;'>No explanation available as no quality tags met the threshold.</div>"
+
+    quality_name = top_quality[0]  # Get the name of the top quality
 
     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:
+Explain why this issue might be classified as a **{quality_name}** issue. Provide a concise explanation, relating it back to the issue description.
 """
-    for quality, _ in top_qualities:
-        prompt += f"- {quality}\n"
-
-
     try:
         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
+        # Format for better readability, directly including the quality name.
+        formatted_explanation = f"<p><b>{quality_name}:</b></p><p>{explanation}</p>"
+        return f"<div style='overflow-y: scroll; max-height: 400px;'>{formatted_explanation}</div>"
     except Exception as e:
         logging.error(f"Error during Llama generation: {e}")
         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
+
+# @spaces.GPU(duration=60)
 def main_interface(text):
     if not text.strip():
         return "<div style='color: red;'>No text provided. Please enter a valid issue description.</div>", "", ""
@@ -172,26 +156,33 @@ def main_interface(text):
     if len(text) < 30:
         return "<div style='color: red;'>Text is less than 30 characters.</div>", "", ""
 
-    device = "cuda" if torch.cuda.is_available() else "cpu" # Keep this for model_prediction
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     results = []
     for model_path, model in models.items():
         quality_name = get_quality_name(model_path)
-        avg_prob = model_prediction(model, text, device)  # Pass the device
-        if avg_prob >= 0.95:
+        avg_prob = model_prediction(model, text, device)
+        if avg_prob >= 0.95:  # Keep *all* results above the threshold
             results.append((quality_name, avg_prob))
         logging.info(f"Model: {model_path}, Quality: {quality_name}, Average Probability: {avg_prob:.3f}")
 
-
     if not results:
-        return "<div style='color: red;'>No recommendation. Prediction probability is below the threshold. </div>", "", ""
+        return "<div style='color: red;'>No recommendation. Prediction probability is below the threshold.</div>", "", ""
+
+    # Sort and get the top result (if any meet the threshold)
+    top_result = sorted(results, key=lambda x: x[1], reverse=True)
+    if top_result:
+        top_quality = top_result[:1] # Select only the top result
+        output_html = render_html_output(top_quality)
+        explanation = generate_explanation(text, top_quality)
+    else: # Handle case no predictions >= 0.95
+        output_html = "<div style='color: red;'>No quality tag met the prediction probability threshold (>= 0.95).</div>"
+        explanation = ""
 
-    top_qualities = sorted(results, key=lambda x: x[1], reverse=True)[:3]
-    output_html = render_html_output(top_qualities)
-    explanation = generate_explanation(text, top_qualities)
 
     return output_html, "", explanation
 
 def render_html_output(top_qualities):
+    #Simplified to show only the top prediction
     styles = """
     <style>
         .quality-container {
@@ -210,25 +201,18 @@ def render_html_output(top_qualities):
             margin-right: 10px;
             box-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
         }
-        .probability {
-            display: block;
-            margin-top: 10px;
-            font-size: 16px;
-            color: #007bff;
-        }
     </style>
     """
-    html_content = ""
-    ranking_labels = ['Top 1 Prediction', 'Top 2 Prediction', 'Top 3 Prediction']
-    top_n = min(len(top_qualities), len(ranking_labels))
-    for i in range(top_n):
-        quality, prob = top_qualities[i]
-        html_content += f"""
-        <div class="quality-container">
-            <span class="ranking">{ranking_labels[i]}</span>
-            <span class="quality-label">{quality}</span>
-        </div>
-        """
+    if not top_qualities: # Handle empty case
+        return styles + "<div class='quality-container'>No Top Prediction</div>"
+
+    quality, _ = top_qualities[0] #We know there is only one
+    html_content = f"""
+    <div class="quality-container">
+        <span class="ranking">Top Prediction</span>
+        <span class="quality-label">{quality}</span>
+    </div>
+    """
     return styles + html_content
 
 example_texts = [
@@ -237,17 +221,48 @@ example_texts = [
     ["There is frequent miscommunication between the development and QA teams regarding feature specifications.\n\nEnvironment: Inter-team meetings\nReproduction: Audit recent communication logs and meeting notes between the teams."],
     ["The service-oriented architecture does not effectively isolate failures, leading to cascading failures across services.\n\nEnvironment: Microservices architecture\nReproduction: Simulate a service failure and observe the impact on other services."]
 ]
-
+# Improved CSS for better layout and appearance
+css = """
+.quality-container {
+    font-family: Arial, sans-serif;
+    text-align: center;
+    margin-top: 20px;
+    padding: 10px;
+    border: 1px solid #ddd; /* Added border */
+    border-radius: 8px; /* Rounded corners */
+    background-color: #f9f9f9; /* Light background */
+}
+.quality-label, .ranking {
+    display: inline-block;
+    padding: 0.5em 1em;
+    font-size: 18px;
+    font-weight: bold;
+    color: white;
+    background-color: #007bff;
+    border-radius: 0.5rem;
+    margin-right: 10px;
+    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
+}
+#explanation {
+  border: 1px solid #ccc;
+  padding: 10px;
+  margin-top: 10px;
+  border-radius: 4px;
+  background-color: #fff; /* White background for explanation */
+  overflow-y: auto;  /* Ensure scrollbar appears if needed */
+}
+"""
 interface = gr.Interface(
     fn=main_interface,
     inputs=gr.Textbox(lines=7, label="Issue Description", placeholder="Enter your issue text here"),
     outputs=[
         gr.HTML(label="Prediction Output"),
         gr.Textbox(label="Predictions", visible=False),
-        gr.HTML(label="Explanation") # Change to gr.HTML
+        gr.HTML(label="Explanation")
     ],
     title="QualityTagger",
     description="This tool classifies text into different quality domains such as Security, Usability, etc., and provides explanations.",
-    examples=example_texts
+    examples=example_texts,
+    css=css  # Apply the CSS
 )
 interface.launch(share=True)