Update app.py

app.py

@@ -57,21 +57,7 @@ models = {path: AutoModelForSequenceClassification.from_pretrained(path) for pat
 def get_quality_name(model_name):
     return quality_mapping.get(model_name.split('/')[-1], "Unknown Quality")
 
-def model_prediction(model, text, device):
-    """Optimized: Using inference_mode and keeping models on GPU during the loop."""
-    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()}
-    
-    with torch.inference_mode(): 
-        outputs = model(**inputs)
-        logits = outputs.logits
-        probs = softmax(logits.cpu().numpy(), axis=1)
-        
-    avg_prob = np.mean(probs[:, 1])
-    return avg_prob
-
+# --- Azure OpenAI Setup ---
 azure_api_key = os.getenv("AZURE_OPENAI_API_KEY")
 
 if not azure_api_key:
@@ -85,9 +71,9 @@ azure_client = AzureOpenAI(
 azure_deployment_name = "gpt-5.4-nano-AM-karthik-prod"
 
 def generate_explanation(issue_text, top_quality):
-    """Generates an explanation using Azure OpenAI."""
+    """Generates an explanation using Azure OpenAI (Pure Markdown output)"""
     if not top_quality:
-        return "<div style='color: red;'>No explanation available as no quality tags met the threshold.</div>"
+        return "**No explanation available** as no quality tags met the threshold."
 
     quality_name = top_quality[0][0]  
 
@@ -110,47 +96,65 @@ Provide your response directly without preamble. Use a clear separation between
                 {"role": "system", "content": "You are an expert software engineering assistant specializing in software quality analysis."},
                 {"role": "user", "content": prompt}
             ],
-            max_completion_tokens=250, # <--- FIXED HERE
+            max_completion_tokens=250, 
             temperature=0.2,
             top_p=0.9
         )
         explanation = response.choices[0].message.content.strip()
-        formatted_explanation = f"<div class='explanation-box'>\n\n### Why this is a {quality_name} issue:\n\n{explanation}\n\n</div>"
+        
+        # FIX 1: Return pure markdown. Do not wrap in HTML <div> tags.
+        formatted_explanation = f"### Why this is a {quality_name} issue:\n\n{explanation}"
         return formatted_explanation
         
     except Exception as e:
         logging.error(f"Error during Azure OpenAI generation: {e}")
-        return "<div style='color: red;'>An error occurred while communicating with the Azure OpenAI service. Please check the logs.</div>"
+        return f"**Error:** An error occurred while communicating with the Azure OpenAI service: {e}"
 
 @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>", "", ""
-
-    if len(text) < 30:
-        return "<div style='color: red;'>Text is less than 30 characters.</div>", "", ""
-
+def run_classification_models(text):
+    """FIX 2: Isolate the GPU work so we aren't holding the GPU while waiting for the network API call."""
     device = "cuda" if torch.cuda.is_available() else "cpu"
     results = []
     
     for model_path, model in models.items():
+        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()}
+        
+        with torch.inference_mode(): 
+            outputs = model(**inputs)
+            logits = outputs.logits
+            probs = softmax(logits.cpu().numpy(), axis=1)
+            
+        avg_prob = np.mean(probs[:, 1])
         quality_name = get_quality_name(model_path)
-        avg_prob = model_prediction(model, text, device)
+        
         if avg_prob >= 0.95: 
             results.append((quality_name, avg_prob))
-        logging.info(f"Model: {model_path}, Quality: {quality_name}, Average Probability: {avg_prob:.3f}")
+            
+    return results
+
+def main_interface(text):
+    if not text.strip():
+        return "<div style='color: red;'>No text provided. Please enter a valid issue description.</div>", "", ""
+
+    if len(text) < 30:
+        return "<div style='color: red;'>Text is less than 30 characters.</div>", "", ""
+
+    # 1. Run the GPU models
+    results = run_classification_models(text)
 
     if not results:
         return "<div style='color: red;'>No recommendation. Prediction probability is below the threshold.</div>", "", ""
 
+    # 2. Process results
     top_result = sorted(results, key=lambda x: x[1], reverse=True)
-    if top_result:
-        top_quality = top_result[:1] 
-        output_html = render_html_output(top_quality)
-        explanation = generate_explanation(text, top_quality)
-    else: 
-        output_html = "<div style='color: red;'>No quality tag met the prediction probability threshold (>= 0.95).</div>"
-        explanation = ""
+    top_quality = top_result[:1] 
+    output_html = render_html_output(top_quality)
+    
+    # 3. Fetch the OpenAI Explanation (Safely running on CPU/Network, off the ZeroGPU queue)
+    explanation = generate_explanation(text, top_quality)
 
     return output_html, "", explanation
 
@@ -174,6 +178,7 @@ example_texts = [
     ["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."]
 ]
 
+# FIX 3: We apply styles directly to Gradio's components using !important overrides
 css = """
 .quality-container {
     font-family: Arial, sans-serif;
@@ -198,23 +203,16 @@ css = """
     box-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
 }
 
+/* Updated: Target the Markdown component natively */
 .explanation-box {
-    border: 1px solid var(--border-color-primary, #ccc);
-    padding: 15px;
-    margin-top: 15px;
-    border-radius: 8px;
-    background-color: var(--background-fill-primary, #fff);
-    color: var(--body-text-color, #000);
-    box-shadow: 0 1px 3px rgba(0,0,0,0.1);
-    line-height: 1.5;
-}
-
-.explanation-box p {
-    margin: 8px 0;
-}
-
-.explanation-box b {
-    color: #007bff;
+    border: 1px solid var(--border-color-primary, #ccc) !important;
+    padding: 15px !important;
+    margin-top: 15px !important;
+    border-radius: 8px !important;
+    background-color: var(--background-fill-primary, #fff) !important;
+    color: var(--body-text-color, #000) !important;
+    box-shadow: 0 1px 3px rgba(0,0,0,0.1) !important;
+    line-height: 1.5 !important;
 }
 
 .dark .quality-container {
@@ -224,20 +222,22 @@ css = """
 }
 
 .dark .explanation-box {
-    background-color: #1f2937;
-    border-color: #374151;
-    color: #f3f4f6;
-    box-shadow: 0 1px 3px rgba(255,255,255,0.05);
+    background-color: #1f2937 !important;
+    border-color: #374151 !important;
+    color: #f3f4f6 !important;
+    box-shadow: 0 1px 3px rgba(255,255,255,0.05) !important;
 }
 """
 
 interface = gr.Interface(
+    # Note: main_interface is no longer @spaces.GPU wrapped!
     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.Markdown(label="Explanation")
+        # FIX 4: Add the CSS class natively to the Markdown block
+        gr.Markdown(label="Explanation", elem_classes="explanation-box")
     ],
     title="QualityTagger",
     description="This tool classifies text into different quality domains such as Security, Usability, Maintainability, Reliability etc., and provides explanations.",