rajat5ranjan commited on
Commit
377f9ae
·
verified ·
1 Parent(s): b7f6cd9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -31
app.py CHANGED
@@ -12,6 +12,8 @@ from langchain.vectorstores import Chroma
12
  import google.generativeai as genai
13
  from langchain_google_genai import GoogleGenerativeAIEmbeddings
14
  from langchain_google_genai import ChatGoogleGenerativeAI
 
 
15
 
16
  GOOGLE_API_KEY=os.environ['GOOGLE_API_KEY']
17
 
@@ -20,46 +22,24 @@ docs = loader.load()
20
 
21
  gemini_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
22
 
23
- # Save to disk
24
- vectorstore = Chroma.from_documents(
25
- documents=docs, # Data
26
- embedding=gemini_embeddings, # Embedding model
27
- persist_directory="./chroma_db" # Directory to save data
28
- )
29
-
30
- vectorstore_disk = Chroma(
31
- persist_directory="./chroma_db", # Directory of db
32
- embedding_function=gemini_embeddings # Embedding model
33
- )
34
- retriever = vectorstore_disk.as_retriever(search_kwargs={"k": 1})
35
- # If there is no environment variable set for the API key, you can pass the API
36
- # key to the parameter `google_api_key` of the `ChatGoogleGenerativeAI` function:
37
- # `google_api_key="key"`.
38
  llm = ChatGoogleGenerativeAI(model="gemini-pro",google_api_key = GOOGLE_API_KEY)
39
 
40
 
41
- llm_prompt_template = """You are an assistant for question-answering tasks.
42
- Use the following context to answer the question.
43
- If you don't know the answer, just say that you don't know.
44
- Use five sentences maximum and keep the answer concise.\n
45
- Question: {question} \nContext: {context} \nAnswer:"""
46
 
47
- llm_prompt = PromptTemplate.from_template(llm_prompt_template)
48
 
 
 
 
49
 
50
- def format_docs(docs):
51
- return "\n\n".join(doc.page_content for doc in docs)
52
 
53
- rag_chain = (
54
- {"context": retriever | format_docs, "question": RunnablePassthrough()}
55
- | llm_prompt
56
- | llm
57
- | StrOutputParser()
58
- )
59
 
60
- prompt = st.text_input("Enter Prompt","What is the best stocks for the next few weeks")
61
 
62
- res = rag_chain.invoke(prompt)
63
  st.write(res)
64
 
65
  # If there is no environment variable set for the API key, you can pass the API
 
12
  import google.generativeai as genai
13
  from langchain_google_genai import GoogleGenerativeAIEmbeddings
14
  from langchain_google_genai import ChatGoogleGenerativeAI
15
+ from langchain.chains.llm import LLMChain
16
+ from langchain.chains import StuffDocumentsChain
17
 
18
  GOOGLE_API_KEY=os.environ['GOOGLE_API_KEY']
19
 
 
22
 
23
  gemini_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
24
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25
  llm = ChatGoogleGenerativeAI(model="gemini-pro",google_api_key = GOOGLE_API_KEY)
26
 
27
 
28
+ # user_prompt = st.text_input("Enter Prompt","What is the best stocks for the next few weeks")
 
 
 
 
29
 
 
30
 
31
+ llm_prompt_template = """You are an assistant for stock market insights.
32
+ Based on the context below
33
+ {context}, Suggest some stocks recommendations"""
34
 
35
+ st.write(llm_prompt_template)
36
+ llm_prompt = PromptTemplate.from_template(llm_prompt_template)
37
 
38
+ llm_chain = LLMChain(llm=llm,prompt=llm_prompt)
39
+ stuff_chain = StuffDocumentsChain(llm_chain=llm_chain,document_variable_name="context")
 
 
 
 
40
 
 
41
 
42
+ res = stuff_chain.invoke(docs)
43
  st.write(res)
44
 
45
  # If there is no environment variable set for the API key, you can pass the API