vivekvar commited on
Commit
7fb98fa
·
verified ·
1 Parent(s): 6b5d076

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -43
app.py CHANGED
@@ -1,38 +1,34 @@
1
  import streamlit as st
2
- from llama_index import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate
3
  from llama_index.llms.huggingface import HuggingFaceInferenceAPI
4
  from dotenv import load_dotenv
5
  from llama_index.embeddings.huggingface import HuggingFaceEmbedding
6
- from llama_index import set_global_service_context
7
  import os
8
  import base64
 
9
 
10
  # Load environment variables
11
  load_dotenv()
12
 
13
- # Configure the Llama index settings for using Hugging Face model
14
- llm = HuggingFaceInferenceAPI(
15
- model_name="bigscience/bloom-7b1", # Use a model available on Hugging Face Inference API
16
- tokenizer_name="bigscience/bloom-7b1",
17
- context_window=2048, # Adjust context window based on the model
18
- api_token=os.getenv("HF_TOKEN"), # Hugging Face API Token
19
  max_new_tokens=512,
20
  generate_kwargs={"temperature": 0.1},
21
  )
22
-
23
- # Set up Hugging Face Embedding model
24
- embed_model = HuggingFaceEmbedding(
25
- model_name="sentence-transformers/all-MiniLM-L6-v2" # Use a suitable embedding model
26
  )
27
 
28
- # Set global service context
29
- service_context = set_global_service_context(llm=llm, embed_model=embed_model)
30
-
31
  # Define the directory for persistent storage and data
32
  PERSIST_DIR = "./db"
33
  DATA_DIR = "data"
34
 
35
- # Ensure data directories exist
36
  os.makedirs(DATA_DIR, exist_ok=True)
37
  os.makedirs(PERSIST_DIR, exist_ok=True)
38
 
@@ -44,50 +40,56 @@ def displayPDF(file):
44
 
45
  def data_ingestion():
46
  documents = SimpleDirectoryReader(DATA_DIR).load_data()
47
- storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
48
- index = VectorStoreIndex.from_documents(documents, service_context=service_context)
49
- index.storage_context.persist()
50
 
51
  def handle_query(query):
52
  storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
53
- index = load_index_from_storage(storage_context, service_context=service_context)
54
  chat_text_qa_msgs = [
55
- (
56
- "user",
57
- """created by vivek created for Neonflake Enterprises OPC Pvt Ltd
58
- Context:
59
- {context}
60
- Question:
61
- {query}
62
- """
63
- )
64
  ]
65
  text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
 
66
  query_engine = index.as_query_engine(text_qa_template=text_qa_template)
 
67
 
68
- try:
69
- answer = query_engine.query({"context": "Extracted context from PDF", "query": query})
70
- return answer.get('response', "Sorry, no answer found.")
71
- except Exception as e:
72
- return f"An error occurred: {str(e)}"
 
 
73
 
74
  # Streamlit app initialization
75
- st.title("Chat with your PDF 📄")
76
  st.markdown("Built by [vivek](https://github.com/saravivek-cyber)")
 
77
 
78
  if 'messages' not in st.session_state:
79
  st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload a PDF and ask me anything about its content.'}]
80
 
81
  with st.sidebar:
82
  st.title("Menu:")
83
- uploaded_file = st.file_uploader("Upload your PDF File")
84
- if uploaded_file:
85
- filepath = os.path.join(DATA_DIR, "saved_pdf.pdf")
86
- with open(filepath, "wb") as f:
87
- f.write(uploaded_file.getbuffer())
88
- st.success("File uploaded successfully. Processing...")
89
- data_ingestion()
90
- st.success("Data ingestion completed.")
 
91
 
92
  user_prompt = st.chat_input("Ask me anything about the content of the PDF:")
93
  if user_prompt:
 
1
  import streamlit as st
2
+ from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate
3
  from llama_index.llms.huggingface import HuggingFaceInferenceAPI
4
  from dotenv import load_dotenv
5
  from llama_index.embeddings.huggingface import HuggingFaceEmbedding
6
+ from llama_index.core import Settings
7
  import os
8
  import base64
9
+ import altair as alt
10
 
11
  # Load environment variables
12
  load_dotenv()
13
 
14
+ # Configure the Llama index settings
15
+ Settings.llm = HuggingFaceInferenceAPI(
16
+ model_name="google/gemma-1.1-7b-it",
17
+ tokenizer_name="google/gemma-1.1-7b-it",
18
+ context_window=3000,
19
+ token=os.getenv("HF_TOKEN"),
20
  max_new_tokens=512,
21
  generate_kwargs={"temperature": 0.1},
22
  )
23
+ Settings.embed_model = HuggingFaceEmbedding(
24
+ model_name="BAAI/bge-small-en-v1.5"
 
 
25
  )
26
 
 
 
 
27
  # Define the directory for persistent storage and data
28
  PERSIST_DIR = "./db"
29
  DATA_DIR = "data"
30
 
31
+ # Ensure data directory exists
32
  os.makedirs(DATA_DIR, exist_ok=True)
33
  os.makedirs(PERSIST_DIR, exist_ok=True)
34
 
 
40
 
41
  def data_ingestion():
42
  documents = SimpleDirectoryReader(DATA_DIR).load_data()
43
+ storage_context = StorageContext.from_defaults()
44
+ index = VectorStoreIndex.from_documents(documents)
45
+ index.storage_context.persist(persist_dir=PERSIST_DIR)
46
 
47
  def handle_query(query):
48
  storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
49
+ index = load_index_from_storage(storage_context)
50
  chat_text_qa_msgs = [
51
+ (
52
+ "user",
53
+ """created by vivek created for Neonflake Enterprises OPC Pvt Ltd
54
+ Context:
55
+ {context_str}
56
+ Question:
57
+ {query_str}
58
+ """
59
+ )
60
  ]
61
  text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
62
+
63
  query_engine = index.as_query_engine(text_qa_template=text_qa_template)
64
+ answer = query_engine.query(query)
65
 
66
+ if hasattr(answer, 'response'):
67
+ return answer.response
68
+ elif isinstance(answer, dict) and 'response' in answer:
69
+ return answer['response']
70
+ else:
71
+ return "Sorry, I couldn't find an answer."
72
+
73
 
74
  # Streamlit app initialization
75
+ st.title("Chat with your PDF📄")
76
  st.markdown("Built by [vivek](https://github.com/saravivek-cyber)")
77
+ st.markdown("chat here")
78
 
79
  if 'messages' not in st.session_state:
80
  st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload a PDF and ask me anything about its content.'}]
81
 
82
  with st.sidebar:
83
  st.title("Menu:")
84
+ uploaded_file = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button")
85
+ if st.button("Submit & Process"):
86
+ with st.spinner("Processing..."):
87
+ filepath = "data/saved_pdf.pdf"
88
+ with open(filepath, "wb") as f:
89
+ f.write(uploaded_file.getbuffer())
90
+ # displayPDF(filepath) # Display the uploaded PDF
91
+ data_ingestion() # Process PDF every time new file is uploaded
92
+ st.success("Done")
93
 
94
  user_prompt = st.chat_input("Ask me anything about the content of the PDF:")
95
  if user_prompt: