68 lines
2.6 KiB
Python
68 lines
2.6 KiB
Python
# main.py
|
|
from pocketflow import BatchNode, Flow
|
|
import collections
|
|
from utils import call_llm
|
|
import yaml
|
|
class MajorityVoteNode(BatchNode):
|
|
def prep(self, shared):
|
|
question = shared.get("question", "(No question provided)")
|
|
attempts_count = shared.get("num_tries", 3)
|
|
return [question for _ in range(attempts_count)]
|
|
|
|
def exec(self, single_question: str):
|
|
prompt = f"""
|
|
You are a helpful assistant. Please answer the user's question below.
|
|
Question: {single_question}
|
|
|
|
Return strictly using the following YAML structure:
|
|
```yaml
|
|
thinking: |
|
|
(Your thinking process here)
|
|
answer: 0.123 # Final answer as a decimal with 3 decimal places
|
|
```"""
|
|
raw_response = call_llm(prompt)
|
|
yaml_part = raw_response.split("```yaml")[1].split("```")[0].strip()
|
|
parsed = yaml.safe_load(yaml_part)
|
|
|
|
# Validate we have at least 'answer' field
|
|
if not isinstance(parsed, dict) or 'answer' not in parsed:
|
|
raise RuntimeError(f"Missing 'answer' in YAML: {parsed}")
|
|
|
|
# Return only the 'answer' field for the majority vote.
|
|
return str(parsed['answer'])
|
|
|
|
def exec_fallback(self, prep_res, exc):
|
|
return None
|
|
|
|
def post(self, shared, prep_res, exec_res_list):
|
|
# Count frequency for non-None answers
|
|
exec_res_list = [res for res in exec_res_list if res is not None]
|
|
counter = collections.Counter(exec_res_list)
|
|
best_answer, freq = counter.most_common(1)[0]
|
|
|
|
# Store final
|
|
shared["majority_answer"] = best_answer
|
|
|
|
print("========================")
|
|
print("All structured answers:", exec_res_list)
|
|
print("Majority vote =>", best_answer)
|
|
print("Frequency =>", freq)
|
|
print("========================")
|
|
|
|
# End the flow
|
|
return "end"
|
|
|
|
if __name__ == "__main__":
|
|
shared = {
|
|
"question": """You work at a shoe factory. In front of you, there are three pairs of shoes (six individual shoes) with the following sizes: two size 4s, two size 5s, and two size 6s. The factory defines an "acceptable pair" as two shoes that differ in size by a maximum of one size (e.g., a size 5 and a size 6 would be an acceptable pair). If you close your eyes and randomly pick three pairs of shoes without replacement, what is the probability that you end up drawing three acceptable pairs?""",
|
|
"num_tries": 5
|
|
}
|
|
|
|
majority_node = MajorityVoteNode()
|
|
flow = Flow(start=majority_node)
|
|
flow.run(shared)
|
|
|
|
print("\n=== Final Answer ===")
|
|
print(shared["majority_answer"])
|
|
print("====================")
|