Spaces:
Runtime error
Runtime error
File size: 5,273 Bytes
8e70e09 |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import yaml
import fitz
import torch
import gradio as gr
from PIL import Image
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain.document_loaders import PyPDFLoader
from langchain.prompts import PromptTemplate
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import spaces
from langchain_text_splitters import CharacterTextSplitter, RecursiveCharacterTextSplitter
class RAGbot:
def __init__(self, config_path="config.yaml"):
self.processed = False
self.page = 0
self.chat_history = []
self.prompt = None
self.documents = None
self.embeddings = None
self.vectordb = None
self.tokenizer = None
self.model = None
self.pipeline = None
self.chain = None
self.chunk_size = 512
self.overlap_percentage = 50
self.max_chunks_in_context = 2
self.current_context = None
self.model_temperatue = 0.5
self.format_seperator = "\n\n--\n\n"
self.pipe = None
with open(config_path, "r") as file:
config = yaml.safe_load(file)
self.model_embeddings = config["modelEmbeddings"]
self.auto_tokenizer = config["autoTokenizer"]
self.auto_model_for_causal_lm = config["autoModelForCausalLM"]
def load_embeddings(self):
self.embeddings = HuggingFaceEmbeddings(model_name=self.model_embeddings)
print("Embedding model loaded")
def load_vectordb(self):
overlap = int((self.overlap_percentage / 100) * self.chunk_size)
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=self.chunk_size,
chunk_overlap=overlap,
length_function=len,
add_start_index=True,
)
docs = text_splitter.split_documents(self.documents)
self.vectordb = Chroma.from_documents(docs, self.embeddings)
print("Vector store created")
@spaces.GPU
def load_tokenizer(self):
self.tokenizer = AutoTokenizer.from_pretrained(self.auto_tokenizer)
@spaces.GPU
def create_organic_pipeline(self):
self.pipe = pipeline(
"text-generation",
model=self.auto_model_for_causal_lm,
model_kwargs={"torch_dtype": torch.bfloat16},
device="cuda",
)
print("Model pipeline loaded")
def get_organic_context(self, query):
documents = self.vectordb.similarity_search_with_relevance_scores(query, k=self.max_chunks_in_context)
context = self.format_seperator.join([doc.page_content for doc, score in documents])
self.current_context = context
print("Context Ready")
print(self.current_context)
@spaces.GPU
def create_organic_response(self, history, query):
self.get_organic_context(query)
messages = [
{"role": "system", "content": "From the context given below, answer the user's question\n" + self.current_context},
{"role": "user", "content": query},
]
prompt = self.pipe.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
temp = 0.1
outputs = self.pipe(
prompt,
max_new_tokens=1024,
do_sample=True,
temperature=temp,
top_p=0.9,
)
print(outputs)
return outputs[0]["generated_text"][len(prompt):]
def process_file(self, file):
self.documents = PyPDFLoader(file.name).load()
self.load_embeddings()
self.load_vectordb()
self.create_organic_pipeline()
@spaces.GPU
def generate_response(self, history, query, file, chunk_size, chunk_overlap_percentage, model_temperature, max_chunks_in_context):
self.chunk_size = chunk_size
self.overlap_percentage = chunk_overlap_percentage
self.model_temperatue = model_temperature
self.max_chunks_in_context = max_chunks_in_context
if not query:
raise gr.Error(message='Submit a question')
if not file:
raise gr.Error(message='Upload a PDF')
if not self.processed:
self.process_file(file)
self.processed = True
result = self.create_organic_response(history="", query=query)
for char in result:
history[-1][-1] += char
return history, ""
def render_file(self, file, chunk_size, chunk_overlap_percentage, model_temperature, max_chunks_in_context):
print(chunk_size)
doc = fitz.open(file.name)
page = doc[self.page]
self.chunk_size = chunk_size
self.overlap_percentage = chunk_overlap_percentage
self.model_temperatue = model_temperature
self.max_chunks_in_context = max_chunks_in_context
pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72))
image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
return image
def add_text(self, history, text):
if not text:
raise gr.Error('Enter text')
history.append((text, ''))
return history
|