Update app.py

app.py

@@ -118,9 +118,14 @@ def llama_generate(
             eos_token_id=llama_tokenizer.eos_token_id,
 
         )
-    output_text = llama_tokenizer.decode(generate_ids[0], skip_special_tokens=True)
+        
+    # Extract only the newly generated tokens
+    input_length = inputs.input_ids.shape[1]
+    generated_tokens = generate_ids[0][input_length:]
+    
+    output_text = llama_tokenizer.decode(generated_tokens, skip_special_tokens=True)
     torch.cuda.empty_cache()
-    return output_text
+    return output_text.strip()
 
 
 def generate_explanation(issue_text, top_quality):
@@ -130,30 +135,23 @@ def generate_explanation(issue_text, top_quality):
 
     quality_name = top_quality[0][0]  # Get the name of the top quality
 
-    prompt = f""" Analyze the following issue description based on the quality dimension: **{quality_name}**.
+    prompt = f"""Analyze the following issue description based on the quality dimension: {quality_name}.
                 
-                Issue Description:
-                ---
-                {issue_text}
-                ---
+Issue Description:
+---
+{issue_text}
+---
                 
-                1. **Justification**: Briefly explain why this issue fails or relates to **{quality_name}**. 
-                2. **Improved Version**: Suggest how to rewrite the issue description to better meet this quality standard (e.g., making it more clear, concise, or actionable).
+1. **Justification**: Briefly explain why this issue fails or relates to {quality_name}. 
+2. **Improved Version**: Suggest how to rewrite the issue description to better meet this quality standard (e.g., making it more clear, concise, or actionable).
                 
-                Provide your response directly without preamble. Use a clear separation between the justification and the rewrite.
-                """
+Provide your response directly without preamble. Use a clear separation between the justification and the rewrite."""
+    
     try:
         explanation = llama_generate(prompt)
-        # Extract only the model's explanation, removing any prompt repetition
-        # This typically removes any preamble like "Here's why this is a [quality] issue:"
-        cleaned_explanation = explanation.split("---")[-1].strip()
-        if cleaned_explanation.lower().startswith(quality_name.lower()):
-            cleaned_explanation = cleaned_explanation[len(quality_name):].strip()
-        if cleaned_explanation.startswith(":"):
-            cleaned_explanation = cleaned_explanation[1:].strip()
-            
-        # Format for better readability
-        formatted_explanation = f"<div class='explanation-box'><p><b>Why this is a {quality_name} issue:</b></p><p>{cleaned_explanation}</p></div>"
+        
+        # Format for better readability. Using linebreaks helps Gradio's Markdown component parse it correctly inside the HTML block.
+        formatted_explanation = f"<div class='explanation-box'>\n\n### Why this is a {quality_name} issue:\n\n{explanation}\n\n</div>"
         return formatted_explanation
     except Exception as e:
         logging.error(f"Error during Llama generation: {e}")
@@ -280,9 +278,9 @@ interface = gr.Interface(
         gr.Markdown(label="Explanation")
     ],
     title="QualityTagger",
-    description="This tool classifies text into different quality domains such as Security, Usability,Mantainability, Reliability etc., and provides explanations.",
+    description="This tool classifies text into different quality domains such as Security, Usability, Mantainability, Reliability etc., and provides explanations.",
     examples=example_texts,
     css=css,
     cache_examples=False# Apply the CSS
 )
-interface.launch( )
+interface.launch()