Update app.py

app.py

@@ -1,74 +1,26 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from playwright.sync_api import sync_playwright
-from flax import linen as nn
-from jax import random
-import jax
-import jax.numpy as jnp
+import json
 
-# Define LLaVA model parameters
+# Define model and inference parameters
 MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
 MAX_LENGTH = 512
 NUM_BEAMS = 5
-
-# Define Flax model for action generation
-class ActionModel(nn.Module):
-    vocab_size: int
-    hidden_size: int
-    num_layers: int
-
-    def setup(self):
-        self.embedding = nn.Embed(self.vocab_size, self.hidden_size)
-        self.lstm = nn.LSTM(self.hidden_size, self.hidden_size, num_layers=self.num_layers)
-        self.dense = nn.Dense(self.vocab_size)
-
-    def __call__(self, inputs, init_state):
-        embedded = self.embedding(inputs)
-        output, new_state = self.lstm(embedded, init_state)
-        logits = self.dense(output)
-        return logits, new_state
-
-# Initialize Flax model and get its initial state
-vocab_size = 50257  # Adjust this if needed for Zephyr-7b-beta
-hidden_size = 1024
-num_layers = 2
-key = random.PRNGKey(0)
-model = ActionModel(vocab_size, hidden_size, num_layers)
-init_state = model.lstm.initialize_carry(key, (1, hidden_size))
+TEMPERATURE = 0.7
 
 # Function to generate actions using Zephyr-7b-beta model
-def generate_actions(input_text, browser, page):
-    # Load Zephyr-7b-beta model
+def generate_actions(input_text):
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-    # Prepare input for Zephyr-7b-beta
-    inputs = tokenizer(input_text, return_tensors="pt")
-    inputs = inputs.to(model.device)
-
-    # Generate response (use pipeline for Zephyr-7b-beta)
+    # Use pipeline for text generation
     generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
-    outputs = generator(input_text, max_length=MAX_LENGTH, num_beams=NUM_BEAMS, temperature=0.7)
+    outputs = generator(input_text, max_length=MAX_LENGTH, num_beams=NUM_BEAMS, temperature=TEMPERATURE)
 
-    # Decode response and extract actions
     response = outputs[0]['generated_text']
     actions = response.split("\n")
-    
-    # Perform actions
-    for action in actions:
-        if "open website" in action:
-            website = action.split(" ")[-1]
-            page.goto(website)
-        elif "click" in action:
-            selector = action.split(" ")[-1]
-            page.click(selector)
-        elif "type" in action:
-            text = action.split(" ")[-1]
-            page.type(text)
-        elif "submit" in action:
-            page.press("Enter")
-        else:
-            print(f"Action not recognized: {action}")
+    return actions
 
 # Function to initialize browser and page
 def initialize_browser():
@@ -81,7 +33,23 @@ def initialize_browser():
 def run_agent(input_text):
     with sync_playwright() as p:
         browser, page = initialize_browser()
-        generate_actions(input_text, browser, page)
+        actions = generate_actions(input_text)
+
+        for action in actions:
+            if "open website" in action:
+                website = action.split(" ")[-1]
+                page.goto(website)
+            elif "click" in action:
+                selector = action.split(" ")[-1]
+                page.click(selector)
+            elif "type" in action:
+                text = action.split(" ")[-1]
+                page.type(text)
+            elif "submit" in action:
+                page.press("Enter")
+            else:
+                print(f"Action not recognized: {action}")
+
         return f"Successfully executed actions based on: {input_text}"
 
 iface = gr.Interface(