|
import gradio as gr |
|
from langchain.document_loaders import PDFMinerLoader, PyMuPDFLoader |
|
from langchain.text_splitter import CharacterTextSplitter |
|
import chromadb |
|
import chromadb.config |
|
from chromadb.config import Settings |
|
from transformers import T5ForConditionalGeneration, AutoTokenizer |
|
import torch |
|
import gradio as gr |
|
import uuid |
|
from sentence_transformers import SentenceTransformer |
|
|
|
model_name = 'google/flan-t5-base' |
|
model = T5ForConditionalGeneration.from_pretrained(model_name, device_map='auto', offload_folder="offload") |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
print('flan read') |
|
|
|
|
|
ST_name = 'sentence-transformers/sentence-t5-base' |
|
st_model = SentenceTransformer(ST_name) |
|
print('sentence read') |
|
|
|
|
|
def get_context(query_text): |
|
query_emb = st_model.encode(query_text) |
|
query_response = collection.query(query_embeddings=query_emb.tolist(), n_results=4) |
|
context = query_response['documents'][0][0] |
|
context = context.replace('\n', ' ').replace(' ', ' ') |
|
return context |
|
|
|
def local_query(query, context): |
|
t5query = """Using the available context, please answer the question. |
|
If you aren't sure please say i don't know. |
|
Context: {} |
|
Question: {} |
|
""".format(context, query) |
|
|
|
inputs = tokenizer(t5query, return_tensors="pt") |
|
outputs = model.generate(**inputs, max_new_tokens=20) |
|
|
|
return tokenizer.batch_decode(outputs, skip_special_tokens=True) |
|
|
|
def run_query(query): |
|
context = get_context(query) |
|
result = local_query(query, context) |
|
return result |
|
|
|
|
|
|
|
def upload_pdf(file): |
|
|
|
file_name = file.name |
|
pdf_filename = os.path.basename(file_path) |
|
|
|
|
|
loader = PDFMinerLoader(pdf_filename) |
|
doc = loader.load() |
|
|
|
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0) |
|
texts = text_splitter.split_documents(doc) |
|
|
|
texts = [i.page_content for i in texts] |
|
|
|
doc_emb = st_model.encode(texts) |
|
doc_emb = doc_emb.tolist() |
|
|
|
ids = [str(uuid.uuid1()) for _ in doc_emb] |
|
|
|
client = chromadb.Client() |
|
collection = client.create_collection("test_db") |
|
|
|
collection.add( |
|
embeddings=doc_emb, |
|
documents=texts, |
|
ids=ids |
|
) |
|
|
|
return 'hello' |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
iface = gr.Interface( |
|
fn=upload_pdf, |
|
inputs="file", |
|
outputs="text", |
|
title="PDF File Uploader", |
|
description="Upload a PDF file and get its filename.", |
|
) |
|
|
|
iface.launch() |
|
|
|
|
