Update app.py

app.py

@@ -1,11 +1,12 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModel
+from langchain_groq import ChatGroq
 from langchain_community.document_loaders import WebBaseLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import ChatPromptTemplate
 from langchain.chains import create_retrieval_chain
 from langchain_community.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
 import numpy as np
 import torch
 import time
@@ -14,16 +15,18 @@ import time
 tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
 model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
 
-def embed_text(texts):
-    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
-    with torch.no_grad():
-        embeddings = model(**inputs).last_hidden_state.mean(dim=1)
-    return embeddings.numpy()
+class CustomHuggingFaceEmbeddings(HuggingFaceEmbeddings):
+    def __init__(self):
+        super().__init__(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-def embedding_function(texts):
-    # This function converts texts to embeddings using the Hugging Face model
-    embeddings = embed_text(texts)
-    return embeddings
+    def embed_documents(self, texts):
+        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
+        with torch.no_grad():
+            embeddings = model(**inputs).last_hidden_state.mean(dim=1)
+        return embeddings.numpy()
+
+# Instantiate embeddings class
+embeddings = CustomHuggingFaceEmbeddings()
 
 if "vector" not in st.session_state:
     st.session_state.loader = WebBaseLoader("https://docs.nvidia.com/cuda/")
@@ -32,10 +35,10 @@ if "vector" not in st.session_state:
     st.session_state.text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     documents = st.session_state.text_splitter.split_documents(st.session_state.docs[:50])
 
-    # Create FAISS index using the custom embedding function
+    # Create FAISS index using the custom embeddings class
     st.session_state.vectors = FAISS.from_texts(
         [doc.page_content for doc in documents], 
-        embedding_function
+        embeddings
     )
 
 st.title("ChatGroq Demo")