|
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("-----------") |
|
|
|
|
|
|
|
embeddings = [] |
|
for doc in documents: |
|
embeddings.extend(doc['embeddings']) |
|
|
|
|
|
api_db = Chroma.from_texts(embeddings, 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).", |
|
) |
|
|
|
|
|
|
|
|