|
import os |
|
|
|
import gradio as gr |
|
from langchain.vectorstores import Chroma |
|
from langchain.document_loaders import PyPDFLoader |
|
from langchain.text_splitter import CharacterTextSplitter |
|
from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings |
|
|
|
|
|
inference_api_key = os.environ['HF'] |
|
api_hf_embeddings = HuggingFaceInferenceAPIEmbeddings( |
|
api_key=inference_api_key, |
|
model_name="sentence-transformers/all-MiniLM-l6-v2" |
|
) |
|
|
|
|
|
loader = PyPDFLoader("./new_papers/ALiBi.pdf") |
|
documents = loader.load() |
|
print("-----------") |
|
print(documents[0]) |
|
print("-----------") |
|
|
|
|
|
text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0) |
|
vdocuments = text_splitter.split_documents(documents) |
|
|
|
|
|
|
|
print("Length of documents: %s" % len(documents)) |
|
|
|
print("Length of vdocuments: %s", len(vdocuments)) |
|
|
|
|
|
if vdocuments and 'embeddings' in vdocuments[0]: |
|
first_document_embeddings = vdocuments[0]['embeddings'] |
|
print("Length of embeddings for the first document: {}".format(len(first_document_embeddings))) |
|
|
|
|
|
|
|
api_db = Chroma.from_documents(vdocuments, api_hf_embeddings, collection_name="api-collection") |
|
|
|
|
|
def pdf_retrieval(query): |
|
|
|
response = api_db.similarity_search(query) |
|
return response |
|
|
|
|
|
|
|
api_tool = gr.Interface( |
|
fn=pdf_retrieval, |
|
inputs=[gr.Textbox()], |
|
outputs=gr.Textbox(), |
|
live=True, |
|
title="API PDF Retrieval Tool", |
|
description="This tool indexes PDF documents and retrieves relevant answers based on a given query (HF Inference API Embeddings).", |
|
) |
|
|
|
|
|
api_tool.launch() |
|
|