Update app.py

app.py

@@ -3,17 +3,15 @@ import streamlit as st
 import pandas as pd
 import sqlite3
 import openai
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline  # Using Hugging Face pipeline for memory-efficient loading
 from langchain import OpenAI
-from langchain.agents import create_sql_agent
-from langchain.sql_database import SQLDatabase
-from langchain.chains import RetrievalQA
-from langchain.document_loaders import CSVLoader
-from langchain.vectorstores import FAISS
-from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain_community.agent_toolkits.sql.base import create_sql_agent
+from langchain_community.utilities import SQLDatabase
+from langchain_community.document_loaders import CSVLoader
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import OpenAIEmbeddings
 import sqlparse
 
-
 # OpenAI API key (ensure it is securely stored)
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
@@ -37,9 +35,8 @@ data.to_sql(table_name, conn, index=False, if_exists='replace')
 # SQL table metadata (for validation and schema)
 valid_columns = list(data.columns)
 
-# Step 3: Use LLaMA for context retrieval (RAG)
-tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
-llama_model = AutoModelForCausalLM.from_pretrained("huggyllama/llama-7b")
+# Step 3: Use a smaller LLaMA for context retrieval (RAG)
+llama_pipeline = pipeline("text-generation", model="huggyllama/llama-2-3b-hf", device=0)  # Use smaller model
 
 # Step 4: Implement RAG with FAISS for vectorized document retrieval
 embeddings = OpenAIEmbeddings()  # You can use other embeddings if preferred
@@ -50,8 +47,6 @@ documents = loader.load()
 vector_store = FAISS.from_documents(documents, embeddings)
 retriever = vector_store.as_retriever()
 
-rag_chain = RetrievalQA.from_chain_type(llama_model, retriever=retriever)
-
 # Step 5: OpenAI for SQL query generation based on user prompt and context
 openai_llm = OpenAI(temperature=0)
 db = SQLDatabase.from_uri('sqlite:///:memory:')  # Create an SQLite database for LangChain
@@ -77,7 +72,7 @@ user_prompt = st.text_input("Enter your natural language prompt:")
 if user_prompt:
     try:
         # Step 9: Retrieve relevant context using LLaMA RAG
-        rag_result = rag_chain.run(user_prompt)
+        rag_result = llama_pipeline(user_prompt, max_length=200)
         st.write(f"Retrieved Context from LLaMA RAG: {rag_result}")
 
         # Step 10: Generate SQL query with OpenAI based on user prompt and retrieved context