Update app.py

app.py

@@ -4,8 +4,12 @@ import re
 import numpy as np
 import pandas as pd
 import os
- 
-# Set seed for reproducibility
+import torch
+
+# Check if CUDA (GPU) is available
+print("CUDA available:", torch.cuda.is_available())
+
+# Set a seed for reproducibility
 set_seed(42)
 
 # Define the six premium generation models:
@@ -33,30 +37,32 @@ grammar_model_names = [
     "hassaanik/grammar-correction-model"
 ]
 
-# Function to load generation pipelines on demand
+# Function to load generation pipelines, specifying GPU if available.
 def load_generation_pipeline(model_name):
     try:
-        return pipeline("text-generation", model=model_name)
+        # Use device=0 if GPU is available; otherwise, use CPU (device=-1)
+        device = 0 if torch.cuda.is_available() else -1
+        return pipeline("text-generation", model=model_name, device=device)
     except Exception as e:
         print(f"Error loading generation model {model_name}: {e}")
         return None
 
-# Function to load grammar evaluation pipelines on demand
+# Function to load grammar evaluation pipelines.
 def load_grammar_pipeline(model_name):
     try:
-        return pipeline("text2text-generation", model=model_name)
+        device = 0 if torch.cuda.is_available() else -1
+        return pipeline("text2text-generation", model=model_name, device=device)
     except Exception as e:
         print(f"Error loading grammar model {model_name}: {e}")
         return None
 
-# Pre-load grammar evaluators
+# Pre-load grammar evaluators.
 rater_models = []
 for model_name in grammar_model_names:
     p = load_grammar_pipeline(model_name)
     if p is not None:
         rater_models.append(p)
 
-# Utility functions to clean text and check for palindromes
 def clean_text(text):
     return re.sub(r'[^a-zA-Z0-9]', '', text.lower())
 
@@ -64,16 +70,15 @@ def is_palindrome(text):
     cleaned = clean_text(text)
     return cleaned == cleaned[::-1]
 
-# Build prompt with clear instructions to output only the palindrome.
+# Updated prompt: instruct output to contain only the palindrome.
 def build_prompt(lang):
     return (
         f"Instruction: Generate a single original palindrome in {lang}.\n"
         "Output only the palindrome. The palindrome should be a continuous text that reads the same forward and backward.\n"
-        "Do not output any additional text or commentary.\n"
+        "Do not output any additional text, commentary, or the prompt itself.\n"
         "Palindrome: "
     )
 
-# Build prompt for grammar evaluation
 def grammar_prompt(pal, lang):
     return (
         f"Rate from 0 to 100 how grammatically correct this palindrome is in {lang}. "
@@ -81,7 +86,6 @@ def grammar_prompt(pal, lang):
         f'"{pal}"\n'
     )
 
-# Extract numeric score from text output
 def extract_score(text):
     match = re.search(r"\d{1,3}", text)
     if match:
@@ -89,7 +93,7 @@ def extract_score(text):
         return min(max(score, 0), 100)
     return 0
 
-# Main benchmark function - runs all tests and saves CSV automatically.
+# Main benchmark function that runs tests and saves CSV results.
 def run_benchmark_all():
     results = []
     for model_name in premium_models:
@@ -105,7 +109,6 @@ def run_benchmark_all():
             valid = is_palindrome(gen_output)
             cleaned_len = len(clean_text(gen_output))
             
-            # Evaluate grammar using both grammar models
             scores = []
             for rater in rater_models:
                 rprompt = grammar_prompt(gen_output, lang)
@@ -135,12 +138,11 @@ def run_benchmark_all():
     print(f"CSV saved to {os.path.abspath(csv_path)}")
     return gr.Dataframe(df), csv_path
 
-# Build the Gradio UI using a Blocks layout
 with gr.Blocks(title="Premium Model Palindrome Benchmark") as demo:
     gr.Markdown("# Premium Model Palindrome Benchmark")
     gr.Markdown(
-        "This benchmark runs automatically over 6 premium text-generation models across 5 languages "
-        "(English, German, Spanish, French, Portuguese) and saves the results to a CSV file upon completion."
+        "This benchmark runs automatically over 6 premium text-generation models across 5 languages and saves the results "
+        "to a CSV file upon completion."
     )
     with gr.Row():
         run_button = gr.Button("Run All Benchmarks")