Update app.py

app.py

@@ -39,78 +39,47 @@ def load_model(model_name, quantized=False, quantized_model_path=None):
 llama_model, llama_tokenizer = load_model(MODEL_NAME)
 prm_model, _ = load_model(None, quantized=True, quantized_model_path=QUANTIZED_PRM_PATH)
 
-def majority_voting(model, tokenizer, prompt, num_samples=5):
+# Strategies
+def majority_voting(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)
-    }
+    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 max(set(outputs), key=outputs.count)
 
-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 best_of_n(prompt, num_samples=5):
+    scored_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 max(scored_outputs, key=lambda x: x[1])[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 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):
+    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))
+    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 max(results, key=lambda x: x[1])[0]
 
 
 def temperature_sampling(model, tokenizer, prompt, temperature=0.7, num_samples=5):
@@ -135,29 +104,6 @@ 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:
@@ -231,7 +177,7 @@ from datetime import datetime
 
 def calculate_metrics(text):
     return {
-        'token_count': len(text.split()),
+        'token_count': len(text.split()), 
         'char_count': len(text),
         'sentence_count': len([s for s in text.split('.') if s.strip()]),
     }
@@ -255,12 +201,14 @@ def create_token_plot(tokens, strategies):
     return plt
 
 def format_metrics(metrics):
+    print(type(metrics))  # Check if it's a list or dictionary
+    print(metrics)  # Inspect its contents
     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
+- Token Count: {metrics[0]['token_count']}
+- Character Count: {metrics[0]['char_count']}
+- Sentence Count: {metrics[0]['sentence_count']}
+- Generation Time: {metrics[0]['generation_time']:.2f}s
 """
 
 def run_single_strategy(prompt, strategy, num_samples):