Update token_counter.py

token_counter.py

@@ -1,15 +1,51 @@
-from transformers import AutoTokenizer
-from transformers import Tool
+import gradio as gr
+from langchain.document_loaders import TextLoader
+from langchain.vectorstores import Chroma
+from langchain.chains import RetrievalQA
+from langchain.embeddings import HuggingFaceInstructEmbeddings 
+from langchain.agents import Tool
+
+# Initialize the HuggingFaceInstructEmbeddings
+hf = HuggingFaceInstructEmbeddings(
+  model_name="hkunlp/instructor-large",
+  embed_instruction="Represent the document for retrieval: ",
+  query_instruction="Represent the query for retrieval: "
+)
+
+# Example texts for the vector store
+texts=["The meaning of life is to love","The meaning of vacation is to relax","Roses are red.","Hack the planet!"]
+
+# Create a Chroma vector store from the example texts
+db = Chroma.from_texts(texts, hf, collection_name="my-collection")
+
+# Create a RetrievalQA chain
+llm = LLM.from_model("vicuna-13b")  # Replace with the appropriate LLM model
+docsearcher = RetrievalQA.from_chain_type(
+  llm=llm, 
+  chain_type="stuff",  # Replace with the appropriate chain type
+  return_source_documents=False,
+  retriever=db.as_retriever(search_type="similarity", search_kwargs={"k": 1})
+)
+
+class VectorStoreRetrieverTool(Tool):
+    name = "vectorstore_retriever"
+    description = "This tool uses LangChain's RetrievalQA to find relevant answers from a vector store based on a given query."
 
-class TokenCounterTool(Tool):
-    name = "text_generator"
-    description = "This is a tool for counting token used by a prompt. It takes a prompt as input and returns the generated text."
     inputs = ["text"]
     outputs = ["text"]
 
-    def __call__(self, prompt: str):
-       # token = os.environ['hf']
-        tokenizer = AutoTokenizer.from_pretrained("lgaalves/gpt2-dolly")
-        tokens = tokenizer(prompt)["input_ids"]
-        return f"{len(tokens)}"
-        
+    def __call__(self, query: str):
+        # Run the query through the RetrievalQA chain
+        response = docsearcher.run(query)
+        return response
+
+# Create the Gradio interface using the HuggingFaceTool
+tool = gr.Interface(
+    VectorStoreRetrieverTool(),
+    live=True,
+    title="LangChain-Application: Vectorstore-Retriever",
+    description="This tool uses LangChain's RetrievalQA to find relevant answers from a vector store based on a given query.",
+)
+
+# Launch the Gradio interface
+tool.launch()