Update app.py

app.py

@@ -1,104 +1,298 @@
 # Install necessary libraries
-#!pip install transformers accelerate datasets gradio sympy
+#!pip install -q transformers accelerate gguf datasets gradio sympy matplotlib pandas
 
-# Import libraries
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import gradio as gr
-import sympy
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
+import matplotlib.pyplot as plt
+import pandas as pd
 
-# Load Model and Tokenizer
+# Define model paths
 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
-PRM_NAME = "RLHFlow/Llama3.1-8B-PRM"
+QUANTIZED_PRM_PATH = hf_hub_download(
+    repo_id="mradermacher/Llama3.1-8B-PRM-Mistral-Data-GGUF",
+    filename="Llama3.1-8B-PRM-Mistral-Data.Q4_K_S.gguf"
+)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load LLaMA model
-def load_model(model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
-    return model.to(device), tokenizer
+def load_model(model_name, quantized=False, quantized_model_path=None):
+    if quantized:
+        n_gpu_layers = -1 if torch.cuda.is_available() else 0
+        model = Llama(
+            model_path=quantized_model_path,
+            n_ctx=2048,
+            n_batch=512,
+            n_gpu_layers=n_gpu_layers,
+            verbose=False
+        )
+        return model, None
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+        return model, tokenizer
 
+# Load models
 llama_model, llama_tokenizer = load_model(MODEL_NAME)
+prm_model, _ = load_model(None, quantized=True, quantized_model_path=QUANTIZED_PRM_PATH)
 
-# Load Process Reward Model (PRM)
-prm_model, prm_tokenizer = load_model(PRM_NAME)
+def majority_voting(model, tokenizer, prompt, num_samples=5):
+    outputs = []
+    if isinstance(model, Llama):
+        for _ in range(num_samples):
+            output = model(prompt, max_tokens=50, temperature=0.7)
+            outputs.append(output["choices"][0]["text"])
+    else:
+        # Prepare inputs
+        input_ids = tokenizer(prompt, return_tensors="pt", padding=True).input_ids.to(device)
+        
+        for _ in range(num_samples):
+            output = model.generate(
+                input_ids,
+                max_new_tokens=50,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+            outputs.append(tokenizer.decode(output[0], skip_special_tokens=True))
+    
+    return {
+        "outputs": outputs,
+        "final_result": max(set(outputs), key=outputs.count)
+    }
 
-# Strategies
-def majority_voting(prompt, num_samples=5):
+def best_of_n(model, tokenizer, prm_model, prompt, num_samples=5):
+    outputs = []
+    if isinstance(model, Llama):
+        for _ in range(num_samples):
+            output = model(prompt, max_tokens=50, temperature=0.7)
+            response = output["choices"][0]["text"]
+            score = len(response.split())
+            outputs.append((response, score))
+    else:
+        input_ids = tokenizer(prompt, return_tensors="pt", padding=True).input_ids.to(device)
+        
+        for _ in range(num_samples):
+            output = model.generate(
+                input_ids,
+                max_new_tokens=50,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+            response = tokenizer.decode(output[0], skip_special_tokens=True)
+            score = len(response.split())
+            outputs.append((response, score))
+    
+    outputs.sort(key=lambda x: x[1], reverse=True)
+    return {
+        "outputs": outputs,
+        "final_result": outputs[0][0]
+    }
+
+def beam_search(model, tokenizer, prompt, num_beams=5):
+    if isinstance(model, Llama):
+        outputs = []
+        for _ in range(num_beams):
+            output = model(prompt, max_tokens=50, temperature=0.7)
+            outputs.append(output["choices"][0]["text"])
+    else:
+        input_ids = tokenizer(prompt, return_tensors="pt", padding=True).input_ids.to(device)
+        
+        outputs = model.generate(
+            input_ids,
+            max_new_tokens=50,
+            num_beams=num_beams,
+            num_return_sequences=num_beams,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+        outputs = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+    
+    return {
+        "outputs": outputs,
+        "final_result": outputs[0]
+    }
+
+
+def temperature_sampling(model, tokenizer, prompt, temperature=0.7, num_samples=5):
     outputs = []
     for _ in range(num_samples):
-        input_ids = llama_tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-        output = llama_model.generate(input_ids, max_new_tokens=50)
-        outputs.append(llama_tokenizer.decode(output[0], skip_special_tokens=True))
-    # Return the most common result
-    return max(set(outputs), key=outputs.count)
-
-def best_of_n(prompt, num_samples=5):
-    scored_outputs = []
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+        output = model.generate(input_ids, max_new_tokens=50, temperature=temperature)
+        outputs.append(tokenizer.decode(output[0], skip_special_tokens=True))
+    return {
+        "outputs": outputs,
+        "final_result": outputs[0]
+    }
+
+def top_p_sampling(model, tokenizer, prompt, top_p=0.9, num_samples=5):
+    outputs = []
     for _ in range(num_samples):
-        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()
-        scored_outputs.append((response, score))
-    # Return the highest scored response
-    return max(scored_outputs, key=lambda x: x[1])[0]
-
-def beam_search(prompt, num_beams=5):
-    input_ids = llama_tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-    outputs = llama_model.generate(input_ids, max_new_tokens=50, num_beams=num_beams, num_return_sequences=num_beams)
-    return [llama_tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
-
-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]
-
-# Gradio Interface
-def inference(prompt, strategy, num_samples, depth, breadth):
-    if strategy == "Majority Voting":
-        return majority_voting(prompt, num_samples)
-    elif strategy == "Best-of-N":
-        return best_of_n(prompt, num_samples)
-    elif strategy == "Beam Search":
-        return beam_search(prompt, num_samples)
-    elif strategy == "DVTS":
-        return dvts(prompt, depth, breadth)
-    else:
-        return "Invalid Strategy"
-
-gr.Interface(
-    fn=inference,
-    inputs=[
-        gr.Textbox(label="Problem Statement", placeholder="Enter your problem here"),
-        gr.Radio(
-            ["Majority Voting", "Best-of-N", "Beam Search", "DVTS"],
-            label="Inference Strategy",
-        ),
-        gr.Slider(1, 10, step=1, value=5, label="Number of Samples"),
-        gr.Slider(1, 5, step=1, value=3, label="Depth (DVTS Only)"),
-        gr.Slider(1, 5, step=1, value=2, label="Breadth (DVTS Only)"),
-    ],
-    outputs="text",
-    title="Dynamic Inference Toolkit",
-    description="Explore test-time compute scaling strategies with Meta's LLaMA model.",
-).launch()
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+        output = model.generate(input_ids, max_new_tokens=50, top_p=top_p)
+        outputs.append(tokenizer.decode(output[0], skip_special_tokens=True))
+    return {
+        "outputs": outputs,
+        "final_result": outputs[0]
+    }
+
+def custom_strategy(prompt, flow):
+    intermediate_results = []
+    for step in flow:
+        strategy = step.get("strategy")
+        params = step.get("params", {})
+        if strategy == "majority_voting":
+            result = majority_voting(prompt, **params)
+        elif strategy == "best_of_n":
+            result = best_of_n(prompt, **params)
+        elif strategy == "beam_search":
+            result = beam_search(prompt, **params)
+        elif strategy == "top_p_sampling":
+            result = top_p_sampling(prompt, **params)
+        else:
+            continue
+        intermediate_results.append({"strategy": strategy, "result": result})
+        prompt = result["final_result"]
+    return intermediate_results
+
+def compare_strategies(model, tokenizer, prm_model, prompt,  num_samples=5):
+    print("Running comparison...")
+    strategies = {
+        "Majority Voting": majority_voting(model, tokenizer, prompt, num_samples=5),
+        "Best-of-N": best_of_n(model, tokenizer, prm_model, prompt, num_samples=5),
+        "Beam Search": beam_search(model, tokenizer, prompt, num_beams=5)
+        #...
+    }
+
+    plt.figure(figsize=(10, 6))
+    plt.bar(strategies.keys(), [len(s["outputs"]) for s in strategies.values()])
+    plt.title("Strategy Comparison")
+    plt.ylabel("Number of Outputs")
+    plt.xticks(rotation=45)
+    plt.tight_layout()
+    plt.show()
+
+    df = pd.DataFrame.from_dict({
+        strategy: {
+            "Final Result": data["final_result"],
+            "Outputs": data["outputs"]
+        } for strategy, data in strategies.items()
+    }, orient="index")
+
+    return strategies, df
+
+def test_generation():
+    sample_prompt = "Explain the concept of neural networks in simple terms."
+    print("Starting generation test...")
+    strategies_results, results_df = compare_strategies(llama_model, llama_tokenizer, prm_model, sample_prompt, 1)
+    print("\nResults DataFrame:")
+    print(results_df)
+    return strategies_results, results_df
+
+
+    #####
+import gradio as gr
+import pandas as pd
+import json
+
+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 run_single_strategy(prompt, strategy, num_samples):
+    if not prompt:
+        return "Please enter a prompt."
+    
+    strategies = {
+        "Majority Voting": lambda: majority_voting(llama_model, llama_tokenizer, prompt, num_samples),
+        "Best-of-N": lambda: best_of_n(llama_model, llama_tokenizer, prm_model, prompt, num_samples),
+        "Beam Search": lambda: beam_search(llama_model, llama_tokenizer, prompt, num_beams=num_samples)
+    }
+    
+    if strategy not in strategies:
+        return "Invalid strategy selected."
+    
+    result = strategies[strategy]()
+    
+    formatted_output = f"""
+### Final Result:
+{result['final_result']}
+
+### All Outputs:
+{format_outputs(result['outputs'])}
+"""
+    return formatted_output
+
+def run_all_strategies(prompt, num_samples):
+    if not prompt:
+        return "Please enter a prompt."
+    
+    strategies_results, results_df = compare_strategies(
+        llama_model, llama_tokenizer, prm_model, prompt, num_samples
+    )
+    
+    # Format the output for display
+    output_text = "# Results from All Strategies\n\n"
+    for strategy, results in strategies_results.items():
+        output_text += f"""
+## {strategy}
+### Final Result:
+{results['final_result']}
+
+### All Outputs:
+{format_outputs(results['outputs'])}
+
+---
+"""
+    
+    return output_text
+
+# Create the Gradio interface
+with gr.Blocks(title="Text Generation Strategies") as demo:
+    gr.Markdown("# Text Generation Strategies Demo")
+    
+    with gr.Row():
+        with gr.Column():
+            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():
+                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():
+            output_display = gr.Markdown(label="Results")
+    
+    # Set up event handlers
+    single_strategy_btn.click(
+        fn=run_single_strategy,
+        inputs=[prompt_input, strategy_dropdown, num_samples],
+        outputs=output_display
+    )
+    
+    all_strategies_btn.click(
+        fn=run_all_strategies,
+        inputs=[prompt_input, num_samples],
+        outputs=output_display
+    )
+
+# Launch the interface
+if __name__ == "__main__":
+    demo.launch(debug=True)