llama2-index / app.py
rajsecrets0's picture
Update app.py
b74fdef verified
import streamlit as st
import os
from llama_index import VectorStoreIndex, ServiceContext
from llama_index.llms import HuggingFaceLLM
from llama_index.prompts.prompts import SimpleInputPrompt
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from llama_index.embeddings import LangchainEmbedding
# Function to set up the Llama2 model and index
def setup_llama2():
system_prompt = """
You are a Q&A assistant. Your goal is to answer questions as
accurately as possible based on the instructions and context provided.
"""
query_wrapper_prompt = SimpleInputPrompt("{query_str}")
llm = HuggingFaceLLM(
context_window=4096,
max_new_tokens=256,
generate_kwargs={"temperature": 0.0, "do_sample": False},
system_prompt=system_prompt,
query_wrapper_prompt=query_wrapper_prompt,
tokenizer_name="meta-llama/Llama-2-7b-chat-hf",
model_name="meta-llama/Llama-2-7b-chat-hf",
device_map="auto",
model_kwargs={"torch_dtype": torch.float16, "load_in_8bit": True}
)
embed_model = LangchainEmbedding(
HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
)
service_context = ServiceContext.from_defaults(
chunk_size=1024,
llm=llm,
embed_model=embed_model
)
index = VectorStoreIndex(service_context=service_context)
return index
# Streamlit app
def main():
st.title("Q&A Assistant with Llama2 and Hugging Face")
# Upload PDF file
pdf_file = st.file_uploader("Upload a PDF file", type=["pdf"])
if pdf_file is not None:
# Process the uploaded PDF file (you may need to adapt this based on your needs)
# For example, you might want to convert the PDF to text or extract relevant information.
# Here, we assume the content is stored in a variable named 'pdf_content'.
pdf_content = pdf_file.read()
# Set up the Llama2 model and index
index = setup_llama2()
# Section for asking questions
st.header("Ask a Question")
# Query input
query = st.text_input("Enter your question:")
if st.button("Submit Question"):
if query:
# Execute the query
response = index.as_query_engine().query(query)
# Display the response
st.header("Answer")
st.markdown(response)
# Run the Streamlit app
if __name__ == "__main__":
main()