Spaces:
Sleeping
Sleeping
File size: 2,896 Bytes
0901470 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import streamlit as st
import torch
from transformers import BitsAndBytesConfig
# Import llama-index and langchain modules
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings, PromptTemplate
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from langchain.embeddings import HuggingFaceEmbeddings
from llama_index.embeddings.langchain import LangchainEmbedding
# ---------------------------
# Configure your LLM and embeddings
# ---------------------------
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 = PromptTemplate("<|USER|>{query_str}<|ASSISTANT|>")
# Configure BitsAndBytes for quantization
quantization_config = BitsAndBytesConfig(
load_in_8bit=True,
bnb_4bit_compute_dtype=torch.float16
)
# Initialize the HuggingFaceLLM with your model settings
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,
"quantization_config": quantization_config
}
)
# Set up the embedding model using Langchain's HuggingFaceEmbeddings
lc_embed_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
embed_model = LangchainEmbedding(lc_embed_model)
# Apply global settings for llama-index
Settings.llm = llm
Settings.embed_model = embed_model
Settings.chunk_size = 1024
# ---------------------------
# Load documents from repository
# ---------------------------
# The "data" folder should be part of your repository with your documents.
DATA_DIR = "data" # Ensure this folder exists and contains your documents.
try:
documents = SimpleDirectoryReader(DATA_DIR).load_data()
except Exception as e:
st.error(f"Error loading documents from '{DATA_DIR}': {e}")
documents = []
if not documents:
st.warning("No documents found in the data folder. Please add your documents and redeploy.")
else:
# Create the vector store index
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()
# ---------------------------
# Streamlit Interface
# ---------------------------
st.title("LLama Index Q&A Assistant")
user_query = st.text_input("Enter your question:")
if user_query:
with st.spinner("Querying..."):
response = query_engine.query(user_query)
st.markdown("### Response:")
st.write(response)
|