Update app.py

app.py

@@ -22,7 +22,7 @@ from selenium.common.exceptions import (
     StaleElementReferenceException,
 )
 from webdriver_manager.chrome import ChromeDriverManager
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, OpenLlamaForCausalLM, pipeline
 import gradio as gr
 import xml.etree.ElementTree as ET
 import torch
@@ -31,23 +31,19 @@ from mysql.connector import errorcode, pooling
 from dotenv import load_dotenv
 from huggingface_hub import login
 
-# Load model directly
-from transformers import AutoTokenizer, AutoModelForMaskedLM
+model_name = "openlm-research/open_llama_3b_v2"  # Or another OpenLlama variant
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = OpenLlamaForCausalLM.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-mpnet-base-v2", clean_up_tokenization_spaces=True)
-model = AutoModelForMaskedLM.from_pretrained("sentence-transformers/all-mpnet-base-v2")
-
-# Define classifier for zero-shot classification
-classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+openllama_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1  # Use GPU if available
+)
 
-# Define nlp using a simple tokenizer
-from transformers import AutoTokenizer
 nlp = AutoTokenizer.from_pretrained("bert-base-uncased")
 
-
-
-
-
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 if not HUGGINGFACE_TOKEN:
     raise ValueError("HUGGINGFACE_TOKEN is not set in the environment variables.")
@@ -907,7 +903,6 @@ def get_latest_csv() -> str:
         logging.error(f"Error retrieving latest CSV: {e}")
         return None
 
-# Chat Response Function with Dynamic Command Handling
 def respond(
     message: str,
     history: list,
@@ -917,12 +912,8 @@ def respond(
     top_p: float,
 ) -> str:
     """
-    Generates a response using the google/flan-t5-xl model based on the user's message and history.
-    Additionally, handles dynamic commands to interact with individual components.
+    Generates a response using OpenLlamaForCausalLM.
     """
-    if chat_pipeline is None:
-        return "Error: Chat model is not loaded."
-
     try:
         # Check if the message contains a command
         command, params = parse_command(message)
@@ -930,20 +921,20 @@ def respond(
             # Execute the corresponding function
             response = execute_command(command, params)
         else:
-            # Generate a regular response using the model
-            prompt = (
-                f"System: {system_message}\n"
-                f"History: {history}\n"
-                f"User: {message}\n"
-                f"Assistant:"
-            )
-            response = chat_pipeline(
-                prompt,
-                max_length=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                num_return_sequences=1,
-            )[0]["generated_text"]
+            # Generate a regular response using OpenLlama
+        prompt = (
+            f"System: {system_message}\n"
+            f"History: {history}\n"
+            f"User: {message}\n"
+            f"Assistant:"
+        )
+        response = openllama_pipeline(
+            prompt,
+            max_length=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+        )[0]["generated_text"]
+
 
             # Extract the assistant's reply
             response = response.split("Assistant:")[-1].strip()