Update app.py

app.py

@@ -1,8 +1,54 @@
+from mabwiser.mab import MAB, LearningPolicy
+
+def bandit_factory(bandit_type, arms):
+  if bandit_type == "Epsilon Greedy":
+    result = MAB(arms=arms,
+             learning_policy=LearningPolicy.EpsilonGreedy(epsilon=0.3),
+             seed=1234)
+  elif bandit_type == "UCB":
+    result = MAB(arms=arms,
+             learning_policy=LearningPolicy.UCB1(alpha=1),
+             seed=1234)
+  elif bandit_type == "Non-Stationary":
+    result = NSBandit(arms=arms, epsilon=0.3, alpha=0.2)
+  else:
+    raise ValueError("Invalid bandit type")
+
+  result.partial_fit(decisions=arms, rewards=[3]*len(arms))
+  return result
+
+class NSBandit:
+  def __init__(self, arms, epsilon, alpha):
+    self.arms = arms
+    self.epsilon = epsilon
+    self.alpha = alpha
+    self.means = {arm: None for arm in arms}
+    self.t = 0
+
+  def partial_fit(self, decisions, rewards):
+    for arm, reward in zip(decisions, rewards):
+      if self.means[arm] is None:
+        self.means[arm] = reward
+      else:
+        self.means[arm] += self.alpha * (reward - self.means[arm])
+    self.t += 1
+
+
+  def predict(self):
+    nones = [t[0] for t in self.means.items() if t[1] is None]
+    if len(nones) > 0:
+      return random.choice(nones)
+
+    best = max(self.means, key=self.means.get)
+
+    if random.random() < self.epsilon:
+      return random.choice(list(set(self.arms) - {best}))
+    else:
+      return max(self.means, key=self.means.get)
+
 import gradio as gr
 import random
 import pandas as pd
-from helpers import *
-from mabwiser.mab import MAB, LearningPolicy
 
 # Load songs dataset
 #file_path = "/content/drive/My Drive/MIT/RealTime/songs_single_genre.csv"