Update app.py

app.py

@@ -2,11 +2,7 @@ import streamlit as st
 import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
-
-try:
-    from groqflow.groqmodel import GroqModel
-except ImportError as e:
-    st.error(f"Failed to import GroqModel. Please ensure all dependencies are installed correctly. Error: {e}")
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # Configure page
 st.set_page_config(page_title="Data Augmentation App", layout="wide")
@@ -29,14 +25,20 @@ st.title("Data Augmentation and Analysis App")
 st.sidebar.title("Upload Your File")
 st.sidebar.markdown("Supported formats: CSV, Excel")
 
-# Load Groq API key from secrets
-try:
-    groq_api_key = st.secrets["HUGGINGFACE_KEY"]
-    groq_model = GroqModel("llama3-8b-8192", api_key=groq_api_key)
-except KeyError:
-    st.error("API key not found in secrets. Please configure your `HUGGINGFACE_KEY` in Streamlit secrets.")
-except Exception as e:
-    st.error(f"Error initializing GroqModel: {e}")
+# Get the Hugging Face API key from secrets
+hf_api_key = st.secrets.get("HUGGINGFACE_KEY")
+if not hf_api_key:
+    st.error("Hugging Face API key not found in secrets.")
+else:
+    # Initialize the model and tokenizer using the API key
+    try:
+        model_name = "llama3-70b-8192"  # Replace with the correct model name if needed
+        model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_api_key)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_api_key)
+        llm_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        st.success(f"Model {model_name} initialized successfully!")
+    except Exception as e:
+        st.error(f"Error initializing model: {e}")
 
 def load_file(uploaded_file):
     """Load the uploaded file."""
@@ -65,14 +67,14 @@ def generate_graph(data, query):
         st.error(f"Error generating graph: {e}")
 
 def handle_query(data, query):
-    """Handle user query using Groq API."""
+    """Handle user query using the LLM."""
     try:
-        if not groq_model:
-            st.error("GroqModel is not initialized. Check for errors in setup.")
+        if not llm_pipeline:
+            st.error("LLM pipeline is not initialized. Check for errors in setup.")
             return
         prompt = f"Given the dataset: {data.to_dict(orient='records')}, answer the following: {query}"
-        response = groq_model.generate(prompt)
-        st.write("Response:", response)
+        response = llm_pipeline(prompt, max_length=200, num_return_sequences=1)
+        st.write("Response:", response[0]['generated_text'])
     except Exception as e:
         st.error(f"Error in LLM processing: {e}")