Spaces:
Runtime error
Runtime error
File size: 2,633 Bytes
2401907 |
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 |
"""
Main script to run the LLM chatbot using the Zigistry framework
This script performs the following tasks:
1. Perform pre-requisites check
2. Configure LLM and embedding model
3. Data ingestion
4. Query handling
5. Launch the chatbot interface
"""
# --- Importing required libraries ---
from zigistry import constants
from zigistry import pre_requisite
import gradio as gr
from llama_index.core import (
StorageContext,
load_index_from_storage,
VectorStoreIndex,
SimpleDirectoryReader,
ChatPromptTemplate,
Settings,
)
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
# -------------------------------------
# --- Perform pre-requisites and pre-requisite check ---
pre_requisite.performPreRequisites()
if not pre_requisite.CorrectlyCompletedPreRequisites():
exit(1)
# ------------------------------------------------------
# --- Configure LLM and embedding model ---
Settings.llm = HuggingFaceInferenceAPI(
model_name=constants.LLM,
tokenizer_name=constants.LLM,
context_window=3000,
token=constants.HF_TOKEN,
max_new_tokens=512,
generate_kwargs={"temperature": constants.TEMPERATURE},
)
Settings.embed_model = HuggingFaceEmbedding(model_name=constants.EMBEDDER)
# -----------------------------------------
def data_ingestion():
"""
Ingest data from the input files and create an index
"""
documents = SimpleDirectoryReader(input_files=constants.FILES).load_data()
index = VectorStoreIndex.from_documents(documents)
index.storage_context.persist(persist_dir=constants.PERSIST_DIR)
def handle_query(query):
"""
Handle the query and return the response
"""
storage_context = StorageContext.from_defaults(persist_dir=constants.PERSIST_DIR)
index = load_index_from_storage(storage_context)
text_qa_template = ChatPromptTemplate.from_messages(constants.LLM_RESPONSE_FORMAT)
query_engine = index.as_query_engine(text_qa_template=text_qa_template)
answer = query_engine.query(query)
if hasattr(answer, "response"):
return answer.response
if isinstance(answer, dict) and "response" in answer:
return answer["response"]
return "Sorry, I couldn't find an answer."
if __name__ == "__main__":
data_ingestion()
# --- Launch the chatbot interface ---
demo = gr.Interface(
fn=handle_query,
inputs="text",
outputs="text",
title="LLM Chatbot",
flagging_mode="never",
)
demo.launch()
# -------------------------------------
|