veeps commited on
Commit
61f09ee
1 Parent(s): b8f9921

trying rag tutorial

Browse files
Files changed (1) hide show
  1. rag.py +38 -0
rag.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from langchain.chat_models import ChatOpenAI
3
+ from langchain.schema import (
4
+ SystemMessage,
5
+ HumanMessage,
6
+ AIMessage
7
+ )
8
+ from datasets import load_dataset
9
+ from pinecone import Pinecone
10
+
11
+
12
+ dataset = load_dataset(
13
+ "jamescalam/llama-2-arxiv-papers-chunked",
14
+ split="train"
15
+ )
16
+
17
+ os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
18
+
19
+ chat = ChatOpenAI(
20
+ openai_api_key = os.environ["OPENAI_API_KEY"],
21
+ model='gpt-3.5-turbo'
22
+
23
+ )
24
+
25
+
26
+ messages = [
27
+ SystemMessage(content="You are a helpful assistant."),
28
+ HumanMessage(content="Hi AI, how are you today?"),
29
+ AIMessage(content="I'm great thank you. How can I help you?"),
30
+ HumanMessage(content="I'd like to understand string theory.")
31
+ ]
32
+
33
+ res = chat(messages)
34
+
35
+ # add latest AI response to messages
36
+ messages.append(res)
37
+
38
+ # connect to pinecone