Update app.py

app.py

@@ -135,6 +135,29 @@ def top_p_sampling(model, tokenizer, prompt, top_p=0.9, num_samples=5):
         "final_result": outputs[0]
     }
 
+def dvts(prompt, depth=3, breadth=2):
+    """
+    Simplified implementation of DVTS: generates a tree of solutions and evaluates branches using PRM.
+    """
+    results = []
+    for _ in range(breadth):
+        input_ids = llama_tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+        output = llama_model.generate(input_ids, max_new_tokens=50)
+        response = llama_tokenizer.decode(output[0], skip_special_tokens=True)
+        score = prm_model(**prm_tokenizer(response, return_tensors="pt").to(device)).logits.mean().item()
+        results.append((response, score))
+    # Select the top responses and expand them recursively
+    for _ in range(depth - 1):
+        best_responses = sorted(results, key=lambda x: x[1], reverse=True)[:breadth]
+        for response, _ in best_responses:
+            input_ids = llama_tokenizer(response, return_tensors="pt").input_ids.to(device)
+            output = llama_model.generate(input_ids, max_new_tokens=50)
+            extended_response = llama_tokenizer.decode(output[0], skip_special_tokens=True)
+            score = prm_model(**prm_tokenizer(extended_response, return_tensors="pt").to(device)).logits.mean().item()
+            results.append((extended_response, score))
+    # Return the best overall response
+    return max(results, key=lambda x: x[1])[0]
+
 def custom_strategy(prompt, flow):
     intermediate_results = []
     for step in flow:
@@ -190,18 +213,57 @@ def test_generation():
 
 
     #####
-import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
+import matplotlib.pyplot as plt
 import pandas as pd
+import gradio as gr
+import time
 import json
+import numpy as np
+from datetime import datetime
+
+def calculate_metrics(text):
+    return {
+        'token_count': len(text.split()),
+        'char_count': len(text),
+        'sentence_count': len([s for s in text.split('.') if s.strip()]),
+    }
 
-def format_outputs(outputs):
-    if isinstance(outputs, list):
-        return "\n\n".join([f"Output {i+1}: {out}" for i, out in enumerate(outputs)])
-    return outputs
+def create_performance_plot(times, strategies):
+    plt.figure(figsize=(10, 5))
+    plt.bar(strategies, times)
+    plt.title('Generation Time by Strategy')
+    plt.ylabel('Time (seconds)')
+    plt.xticks(rotation=45)
+    plt.tight_layout()
+    return plt
+
+def create_token_plot(tokens, strategies):
+    plt.figure(figsize=(10, 5))
+    plt.bar(strategies, tokens)
+    plt.title('Output Token Count by Strategy')
+    plt.ylabel('Number of Tokens')
+    plt.xticks(rotation=45)
+    plt.tight_layout()
+    return plt
+
+def format_metrics(metrics):
+    return f"""
+### Metrics
+- Token Count: {metrics['token_count']}
+- Character Count: {metrics['char_count']}
+- Sentence Count: {metrics['sentence_count']}
+- Generation Time: {metrics['generation_time']:.2f}s
+"""
 
 def run_single_strategy(prompt, strategy, num_samples):
     if not prompt:
-        return "Please enter a prompt."
+        return "Please enter a prompt.", None, None, None
+    
+    start_time = time.time()
     
     strategies = {
         "Majority Voting": lambda: majority_voting(llama_model, llama_tokenizer, prompt, num_samples),
@@ -210,89 +272,133 @@ def run_single_strategy(prompt, strategy, num_samples):
     }
     
     if strategy not in strategies:
-        return "Invalid strategy selected."
+        return "Invalid strategy selected.", None, None, None
     
     result = strategies[strategy]()
+    generation_time = time.time() - start_time
+    
+    # Calculate metrics
+    metrics = calculate_metrics(result['final_result'])
+    metrics['generation_time'] = generation_time
+    
+    # Create visualizations
+    performance_fig = create_performance_plot([generation_time], [strategy])
+    token_fig = create_token_plot([metrics['token_count']], [strategy])
     
     formatted_output = f"""
-### Final Result:
+# Results for {strategy}
+
+## Final Result
 {result['final_result']}
 
-### All Outputs:
+{format_metrics(metrics)}
+
+## All Outputs
 {format_outputs(result['outputs'])}
+
+## Generation Details
+- Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
+- Number of samples: {num_samples}
+- Model: {MODEL_NAME}
+- Device: {device}
 """
-    return formatted_output
+    
+    return formatted_output, performance_fig, token_fig, metrics
 
 def run_all_strategies(prompt, num_samples):
     if not prompt:
-        return "Please enter a prompt."
+        return "Please enter a prompt.", None, None, None
     
-    strategies_results, results_df = compare_strategies(
-        llama_model, llama_tokenizer, prm_model, prompt, num_samples
-    )
+    all_metrics = {}
+    all_times = []
+    all_tokens = []
+    strategies = ["Majority Voting", "Best-of-N", "Beam Search"]
     
-    # Format the output for display
     output_text = "# Results from All Strategies\n\n"
-    for strategy, results in strategies_results.items():
+    
+    for strategy in strategies:
+        start_time = time.time()
+        result = run_single_strategy(prompt, strategy, num_samples)[0]
+        generation_time = time.time() - start_time
+        
+        metrics = calculate_metrics(result)
+        metrics['generation_time'] = generation_time
+        all_metrics[strategy] = metrics
+        
+        all_times.append(generation_time)
+        all_tokens.append(metrics['token_count'])
+        
         output_text += f"""
 ## {strategy}
-### Final Result:
-{results['final_result']}
-
-### All Outputs:
-{format_outputs(results['outputs'])}
-
+{result}
 ---
 """
     
-    return output_text
+    # Create comparison visualizations
+    performance_fig = create_performance_plot(all_times, strategies)
+    token_fig = create_token_plot(all_tokens, strategies)
+    
+    # Add comparison summary
+    output_text += """
+# Strategy Comparison Summary
+"""
+    for strategy, metrics in all_metrics.items():
+        output_text += f"""
+## {strategy}
+{format_metrics(metrics)}
+"""
+    
+    return output_text, performance_fig, token_fig, all_metrics
 
-# Create the Gradio interface
-with gr.Blocks(title="Text Generation Strategies") as demo:
-    gr.Markdown("# Text Generation Strategies Demo")
+# Create the enhanced Gradio interface
+with gr.Blocks(title="Advanced Text Generation Strategies") as demo:
+    gr.Markdown("# Advanced Text Generation Strategies Demo")
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=2):
             prompt_input = gr.Textbox(
                 label="Enter your prompt",
                 placeholder="Type your prompt here...",
                 lines=3
             )
-            num_samples = gr.Slider(
-                minimum=1,
-                maximum=10,
-                value=5,
-                step=1,
-                label="Number of samples/beams"
-            )
-            
             with gr.Row():
+                num_samples = gr.Slider(
+                    minimum=1,
+                    maximum=10,
+                    value=5,
+                    step=1,
+                    label="Number of samples/beams"
+                )
                 strategy_dropdown = gr.Dropdown(
                     choices=["Majority Voting", "Best-of-N", "Beam Search"],
                     label="Select Strategy",
                     value="Majority Voting"
                 )
-                
+            
             with gr.Row():
                 single_strategy_btn = gr.Button("Run Selected Strategy")
                 all_strategies_btn = gr.Button("Run All Strategies")
         
-        with gr.Column():
+        with gr.Column(scale=3):
             output_display = gr.Markdown(label="Results")
+            with gr.Row():
+                performance_plot = gr.Plot(label="Performance Comparison")
+                token_plot = gr.Plot(label="Token Count Comparison")
+            
+            metrics_display = gr.JSON(label="Detailed Metrics")
     
     # Set up event handlers
     single_strategy_btn.click(
         fn=run_single_strategy,
         inputs=[prompt_input, strategy_dropdown, num_samples],
-        outputs=output_display
+        outputs=[output_display, performance_plot, token_plot, metrics_display]
     )
     
     all_strategies_btn.click(
         fn=run_all_strategies,
         inputs=[prompt_input, num_samples],
-        outputs=output_display
+        outputs=[output_display, performance_plot, token_plot, metrics_display]
     )
 
-# Launch the interface
 if __name__ == "__main__":
     demo.launch(debug=True)