import gradio as gr
import os
import torch
import numpy as np
import random
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM, TextIteratorStreamer
from scipy.special import softmax
import logging
import spaces
from threading import Thread
from collections.abc import Iterator
import csv
# Login to Hugging Face
token = os.getenv("hf_token")
login(token=token)
# Increase CSV field size limit
csv.field_size_limit(1000000)
# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
# Set a seed for reproducibility
seed = 42
np.random.seed(seed)
random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
model_paths = [
'karths/binary_classification_train_port',
'karths/binary_classification_train_perf',
"karths/binary_classification_train_main",
"karths/binary_classification_train_secu",
"karths/binary_classification_train_reli",
"karths/binary_classification_train_usab",
"karths/binary_classification_train_comp"
]
quality_mapping = {
'binary_classification_train_port': 'Portability',
'binary_classification_train_main': 'Maintainability',
'binary_classification_train_secu': 'Security',
'binary_classification_train_reli': 'Reliability',
'binary_classification_train_usab': 'Usability',
'binary_classification_train_perf': 'Performance',
'binary_classification_train_comp': 'Compatibility'
}
# Pre-load models and tokenizer for quality prediction
tokenizer = AutoTokenizer.from_pretrained("distilbert/distilroberta-base")
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)
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.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probs = softmax(logits.cpu().numpy(), axis=1)
avg_prob = np.mean(probs[:, 1])
model.to("cpu")
return avg_prob
# --- Llama 3.2 3B Model Setup ---
LLAMA_MAX_MAX_NEW_TOKENS = 512
LLAMA_DEFAULT_MAX_NEW_TOKENS = 250
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-1B-Instruct"
llama_tokenizer = AutoTokenizer.from_pretrained(llama_model_id)
llama_model = AutoModelForCausalLM.from_pretrained(
llama_model_id,
device_map="auto",
torch_dtype=torch.bfloat16,
)
llama_model.eval()
if llama_tokenizer.pad_token is None:
llama_tokenizer.pad_token = llama_tokenizer.eos_token
def llama_generate(
message: str,
max_new_tokens: int = LLAMA_DEFAULT_MAX_NEW_TOKENS,
temperature: float = 0.2,
top_p: float = 0.9,
top_k: int = 50,
repetition_penalty: float = 1.2,
) -> str:
inputs = llama_tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=LLAMA_MAX_INPUT_TOKEN_LENGTH).to(llama_model.device)
if inputs.input_ids.shape[1] > LLAMA_MAX_INPUT_TOKEN_LENGTH:
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.")
with torch.no_grad():
generate_ids = llama_model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
top_p=top_p,
top_k=top_k,
temperature=temperature,
num_beams=1,
repetition_penalty=repetition_penalty,
pad_token_id=llama_tokenizer.pad_token_id,
eos_token_id=llama_tokenizer.eos_token_id,
)
# 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.strip()
def generate_explanation(issue_text, top_quality):
"""Generates an explanation for the *single* top quality above threshold."""
if not top_quality:
return "No explanation available as no quality tags met the threshold.
"
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}.
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).
Provide your response directly without preamble. Use a clear separation between the justification and the rewrite."""
try:
explanation = llama_generate(prompt)
# Format for better readability. Using linebreaks helps Gradio's Markdown component parse it correctly inside the HTML block.
formatted_explanation = f"\n\n### Why this is a {quality_name} issue:\n\n{explanation}\n\n
"
return formatted_explanation
except Exception as e:
logging.error(f"Error during Llama generation: {e}")
return "An error occurred while generating the explanation.
"
@spaces.GPU(duration=60)
def main_interface(text):
if not text.strip():
return "No text provided. Please enter a valid issue description.
", "", ""
if len(text) < 30:
return "Text is less than 30 characters.
", "", ""
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)
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 "No recommendation. Prediction probability is below the threshold.
", "", ""
# 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 = "No quality tag met the prediction probability threshold (>= 0.95).
"
explanation = ""
return output_html, "", explanation
def render_html_output(top_qualities):
# Cleaned up: Removed the duplicate inline