Spaces:
Sleeping
Sleeping
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() | |