import gradio as gr import os import json from openai import OpenAI import sys # Added for flushing output in case of direct printing # Load sensitive information from environment variables RUNPOD_API_KEY = os.getenv('RUNPOD_API_KEY') RUNPOD_ENDPOINT_ID = os.getenv('RUNPOD_ENDPOINT_ID') # --- Basic Input Validation --- if not RUNPOD_API_KEY: raise ValueError("RunPod API key not found. Please set the RUNPOD_API_KEY environment variable.") if not RUNPOD_ENDPOINT_ID: raise ValueError("RunPod Endpoint ID not found. Please set the RUNPOD_ENDPOINT_ID environment variable.") BASE_URL = f"https://api.runpod.ai/v2/{RUNPOD_ENDPOINT_ID}/openai/v1" MODEL_NAME = "karths/coder_commit_32B" # The specific model hosted on RunPod MAX_TOKENS = 4096 # Max tokens for the model response # --- OpenAI Client Initialization --- client = OpenAI( api_key=RUNPOD_API_KEY, base_url=BASE_URL, ) # --- Gradio App Configuration --- title = "Python Maintainability Refactoring demo" description = """ ## Instructions for Using the Model ### Model Loading Time: - Please allow time for the model on GPU server to initialize if it's starting fresh ("Cold Start"). The response will appear token by token. ### Code Submission: - You can enter or paste your Python code you wish to have refactored, or use the provided example. ### Python Code Constraints: - Keep the code reasonably sized. Large code blocks might face limitations depending on the GPU instance and model constraints. Max response length is set to {} tokens. ### Understanding Changes: - It's important to read the "Changes made" section (if provided by the model) in the refactored code response. This will help in understanding what modifications have been made. ### Usage Recommendation: - Intended for research and evaluation purposes. """.format(MAX_TOKENS) system_prompt = """### Instruction: Refactor the provided Python code to improve its maintainability and efficiency and reduce complexity. Include the refactored code along with comments on the changes made for improving the metrics. ### Input: """ css = """.toast-wrap { display: none !important } """ examples = [ ["""def analyze_sales_data(sales_records): active_sales = filter(lambda record: record['status'] == 'active', sales_records) sales_by_category = {} for record in active_sales: category = record['category'] total_sales = record['units_sold'] * record['price_per_unit'] if category not in sales_by_category: sales_by_category[category] = {'total_sales': 0, 'total_units': 0} sales_by_category[category]['total_sales'] += total_sales sales_by_category[category]['total_units'] += record['units_sold'] average_sales_data = [] for category, data in sales_by_category.items(): average_sales = data['total_sales'] / data['total_units'] if data['total_units'] > 0 else 0 # Avoid division by zero sales_by_category[category]['average_sales'] = average_sales average_sales_data.append((category, average_sales)) average_sales_data.sort(key=lambda x: x[1], reverse=True) for rank, (category, _) in enumerate(average_sales_data, start=1): sales_by_category[category]['rank'] = rank return sales_by_category"""], ["""import pandas as pd import re import ast from code_bert_score import score # Assuming this library is available in the environment import numpy as np def preprocess_code(source_text): def remove_comments_and_docstrings(source_code): # Remove single-line comments source_code = re.sub(r'#.*', '', source_code) # Remove multi-line strings (docstrings) source_code = re.sub(r'(\'\'\'(.*?)\'\'\'|\"\"\"(.*?)\"\"\")', '', source_code, flags=re.DOTALL) return source_code.strip() # Added strip # Pattern to extract code specifically from markdown blocks if present pattern = r"```python\s+(.+?)\s+```" matches = re.findall(pattern, source_text, re.DOTALL) code_to_process = '\n'.join(matches) if matches else source_text cleaned_code = remove_comments_and_docstrings(code_to_process) return cleaned_code def evaluate_dataframe(df): results = {'P': [], 'R': [], 'F1': [], 'F3': []} for index, row in df.iterrows(): try: # Ensure inputs are lists of strings cands = [preprocess_code(str(row['generated_text']))] # Added str() conversion refs = [preprocess_code(str(row['output']))] # Added str() conversion # Ensure code_bert_score.score returns four values score_results = score(cands, refs, lang='python') if len(score_results) == 4: P, R, F1, F3 = score_results results['P'].append(P.item() if hasattr(P, 'item') else P) # Handle potential tensor output results['R'].append(R.item() if hasattr(R, 'item') else R) results['F1'].append(F1.item() if hasattr(F1, 'item') else F1) results['F3'].append(F3.item() if hasattr(F3, 'item') else F3) # Assuming F3 is returned else: print(f"Warning: Unexpected number of return values from score function for row {index}. Got {len(score_results)} values.") for key in results.keys(): results[key].append(np.nan) # Append NaN for unexpected format except Exception as e: print(f"Error processing row {index}: {e}") for key in results.keys(): results[key].append(np.nan) # Use NaN for errors df_metrics = pd.DataFrame(results) return df_metrics def evaluate_dataframe_multiple_runs(df, runs=3): all_results = [] print(f"Starting evaluation for {runs} runs...") for run in range(runs): print(f"Run {run + 1}/{runs}") df_metrics = evaluate_dataframe(df.copy()) # Use a copy to avoid side effects if df is modified all_results.append(df_metrics) print(f"Run {run + 1} completed.") if not all_results: print("No results collected.") return pd.DataFrame(), pd.DataFrame() # Concatenate results and calculate statistics try: concatenated_results = pd.concat(all_results) df_metrics_mean = concatenated_results.groupby(level=0).mean() df_metrics_std = concatenated_results.groupby(level=0).std() print("Mean and standard deviation calculated.") except Exception as e: print(f"Error calculating statistics: {e}") # Return empty DataFrames or handle as appropriate return pd.DataFrame(), pd.DataFrame() return df_metrics_mean, df_metrics_std"""] ] # --- Core Logic (Modified for Streaming) --- def gen_solution_stream(prompt): """ Generates a solution for a given problem prompt by calling the LLM via RunPod and yielding the response chunks as they arrive (streaming). Parameters: - prompt (str): The problem prompt including the system message and user input. Yields: - str: Chunks of the generated solution text. - str: An error message if an exception occurs. """ try: # Call the OpenAI compatible endpoint on RunPod with streaming enabled stream = client.chat.completions.create( model=MODEL_NAME, messages=[{"role": "user", "content": prompt}], temperature=0.1, # Keep temperature low for deterministic refactoring top_p=1.0, max_tokens=MAX_TOKENS, stream=True # Enable streaming ) # Yield content chunks from the stream for chunk in stream: if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content: content = chunk.choices[0].delta.content yield content # Optional: Handle finish reason if needed # if chunk.choices and chunk.choices[0].finish_reason: # print(f"\nStream finished with reason: {chunk.choices[0].finish_reason}") except Exception as e: error_message = f"Error: Could not get streaming response from the model. Details: {str(e)}" print(error_message, file=sys.stderr) # Log error to stderr yield error_message # Yield the error message to be displayed in the UI # --- Gradio Interface Function (Modified for Streaming) --- def predict(message, history): """ Handles the user input, calls the backend model stream, and yields the response chunks. 'history' parameter is required by gr.ChatInterface but might not be used here. """ # Construct the full prompt input_prompt = system_prompt + str(message) # Get the refactored code stream from the backend response_stream = gen_solution_stream(input_prompt) # Yield each chunk received from the stream generator # Gradio's ChatInterface handles accumulating these yields into the chatbot output buffer = "" for chunk in response_stream: buffer += chunk yield buffer # Yield the accumulated buffer to update the UI incrementally # --- Launch Gradio Interface --- # Use gr.ChatInterface for a chat-like experience gr.ChatInterface( predict, # Pass the generator function chatbot=gr.Chatbot(height=500, label="Refactored Code and Explanation", show_copy_button=True), # Added copy button textbox=gr.Textbox(lines=10, label="Python Code", placeholder="Enter or Paste your Python code here..."), title=title, description=description, theme="abidlabs/Lime", # Or choose another theme e.g., gr.themes.Default() examples=examples, cache_examples=False, # Consider enabling caching if examples don't change often submit_btn="Submit Code", retry_btn="Retry", undo_btn="Undo", clear_btn="Clear", css=css # Apply custom CSS if needed ).queue().launch(share=True) # share=True creates a public link (use with caution)