Spaces:
Sleeping
Sleeping
File size: 3,441 Bytes
040f3b6 0901470 040f3b6 0901470 040f3b6 570c187 0901470 570c187 0901470 040f3b6 0901470 570c187 0901470 040f3b6 0901470 040f3b6 0901470 040f3b6 0901470 040f3b6 0901470 040f3b6 0901470 040f3b6 |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import os
import streamlit as st
import torch
from transformers import BitsAndBytesConfig
# Import necessary modules from llama-index and langchain
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings, PromptTemplate
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from langchain.embeddings import HuggingFaceEmbeddings
from llama_index.embeddings.langchain import LangchainEmbedding
# ---------------------------
# Retrieve Hugging Face Token from Environment Variables
# ---------------------------
hf_token = os.getenv("HF_TOKEN")
if hf_token is None:
st.error("Missing Hugging Face token. Please set HF_TOKEN in your Space secrets.")
st.stop()
# ---------------------------
# Configure BitsAndBytes Quantization (only if GPU is available)
# ---------------------------
if torch.cuda.is_available():
quantization_config = BitsAndBytesConfig(
load_in_8bit=True,
bnb_4bit_compute_dtype=torch.float16
)
else:
# If no GPU is available, disable bitsandbytes quantization
quantization_config = None
# ---------------------------
# 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|>")
# Prepare model_kwargs based on whether quantization is enabled
model_kwargs = {"torch_dtype": torch.float16}
if quantization_config is not None:
model_kwargs["quantization_config"] = quantization_config
# Initialize the HuggingFaceLLM with your model settings and authentication token
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=model_kwargs,
)
# 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
# ---------------------------
DATA_DIR = "data" # Ensure this folder exists in your repository 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.")
st.stop()
else:
# Create the vector store index and query engine
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)
|