Advertisement
Guest User

csvLoader

a guest
Oct 18th, 2023
205
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.77 KB | None | 0 0
  1. import os
  2. import sys
  3.  
  4. import openai
  5. from langchain.chains import ConversationalRetrievalChain, RetrievalQA
  6. from langchain.chat_models import ChatOpenAI
  7. from langchain.document_loaders import DirectoryLoader, TextLoader, CSVLoader, UnstructuredCSVLoader
  8. from langchain.embeddings import OpenAIEmbeddings
  9. from langchain.indexes import VectorstoreIndexCreator
  10. from langchain.indexes.vectorstore import VectorStoreIndexWrapper
  11. from langchain.llms import OpenAI
  12. from langchain.vectorstores import Chroma
  13.  
  14. import constants
  15.  
  16. os.environ["OPENAI_API_KEY"] = constants.APIKEY
  17.  
  18. # Enable to save to disk & reuse the model (for repeated queries on the same data)
  19. PERSIST = False
  20.  
  21. query = None
  22. if len(sys.argv) > 1:
  23.     query = sys.argv[1]
  24.  
  25. if PERSIST and os.path.exists("persist"):
  26.     print("Reusing index...\n")
  27.     vectorstore = Chroma(persist_directory="persist", embedding_function=OpenAIEmbeddings())
  28.     index = VectorStoreIndexWrapper(vectorstore=vectorstore)
  29. else:
  30.     # https://python.langchain.com/docs/integrations/document_loaders/csv
  31.     loader = CSVLoader("data/test/net-ex.csv")
  32.     if PERSIST:
  33.         index = VectorstoreIndexCreator(vectorstore_kwargs={"persist_directory": "persist"}).from_loaders([loader])
  34.     else:
  35.         index = VectorstoreIndexCreator().from_loaders([loader])
  36.  
  37. chain = ConversationalRetrievalChain.from_llm(
  38.     llm=ChatOpenAI(model="gpt-3.5-turbo"),
  39.     retriever=index.vectorstore.as_retriever(search_kwargs={"k": 1}),
  40. )
  41.  
  42. chat_history = []
  43. while True:
  44.     if not query:
  45.         query = input("Prompt: ")
  46.     if query in ['quit', 'q', 'exit']:
  47.         sys.exit()
  48.  
  49.     result = chain({"question": query, "chat_history": chat_history})
  50.     print(result['answer'])
  51.  
  52.     chat_history.append((query, result['answer']))
  53.     query = None
  54.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement