Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from playwright.sync_api import sync_playwright
 from flax import linen as nn
 from jax import random
@@ -28,34 +28,30 @@ class ActionModel(nn.Module):
         logits = self.dense(output)
         return logits, new_state
 
-# Initialize Flax model
-vocab_size = 50257
+# 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))
 
-# Function to generate actions using LLaVA model
+# Function to generate actions using Zephyr-7b-beta model
 def generate_actions(input_text, browser, page):
-    # Load LLaVA model
+    # Load Zephyr-7b-beta model
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-    # Prepare input for LLaVA
+    # Prepare input for Zephyr-7b-beta
     inputs = tokenizer(input_text, return_tensors="pt")
     inputs = inputs.to(model.device)
 
-    # Generate response
-    outputs = model.generate(
-        input_ids=inputs.input_ids,
-        max_length=MAX_LENGTH,
-        num_beams=NUM_BEAMS,
-        temperature=0.7,
-    )
+    # Generate response (use pipeline for Zephyr-7b-beta)
+    generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
+    outputs = generator(input_text, max_length=MAX_LENGTH, num_beams=NUM_BEAMS, temperature=0.7)
 
     # Decode response and extract actions
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = outputs[0]['generated_text']
     actions = response.split("\n")
     
     # Perform actions