Spaces:
Running
Running
| ########################################################################################### | |
| # Title: Gradio Interface to LLM-chatbot with dynamic RAG-funcionality and ChromaDB | |
| # Author: Andreas Fischer | |
| # Date: October 10th, 2024 | |
| # Last update: October 26th, 2024 | |
| ########################################################################################## | |
| import os | |
| import torch | |
| from transformers import AutoTokenizer, AutoModel # chromaDB | |
| from datetime import datetime, date #add_doc, | |
| import chromadb #chromaDB | |
| from chromadb import Documents, EmbeddingFunction, Embeddings #chromaDB | |
| from chromadb.utils import embedding_functions #chromaDB | |
| import ocrmypdf #convertPDF | |
| from pypdf import PdfReader #convertPDF | |
| import re #format_prompt | |
| import gradio as gr # multimodal_response | |
| from huggingface_hub import InferenceClient # multimodal_response | |
| import json # multimodal_response (on-prem) | |
| import requests # multimodal_response (on-prem) | |
| #--------------------------------------------------- | |
| # Specify models for text generation and embeddings | |
| #--------------------------------------------------- | |
| myModel="mistralai/Mixtral-8x7b-instruct-v0.1" | |
| #myModel="meta-llama/Llama-3.1-8B-Instruct" | |
| #myModel="QuantFactory/gemma-2-9b-it-SimPO-GGUF" | |
| #myModel="bartowski/gemma-2-9b-it-GGUF" | |
| #mod=myModel | |
| #tok=AutoTokenizer.from_pretrained(mod) #,token="hf_...") | |
| #cha=[{"role":"system","content":"A"},{"role":"user","content":"B"},{"role":"assistant","content":"C"}] | |
| #cha=[{"role":"user","content":"U1"},{"role":"assistant","content":"A1"},{"role":"user","content":"U2"},{"role":"assistant","content":"A2"}] | |
| #res=tok.apply_chat_template(cha) | |
| #print(tok.decode(res)) | |
| if("GGUF" in myModel): # start Llama-cpp-server for GGUF-models on premises: | |
| #modelPath="/home/af/gguf/models/bartowski/gemma-2-9b-it-GGUF/gemma-2-9b-it-Q4_K_M.gguf" | |
| modelPath="/home/af/gguf/models/QuantFactory/gemma-2-9b-it-SimPO-GGUF/gemma-2-9b-it-SimPO.Q4_K_M.gguf" | |
| if(os.path.exists(modelPath)==False): | |
| #url="https://huggingface.co/bartowski/gemma-2-9b-it-GGUF/resolve/main/gemma-2-9b-it-Q4_K_M.gguf?download=true" | |
| url="https://huggingface.co/QuantFactory/gemma-2-9b-it-SimPO-GGUF/resolve/main/gemma-2-9b-it-SimPO.Q4_K_M.gguf?download=true" | |
| response = requests.get(url) | |
| with open("./model.gguf", mode="wb") as file: | |
| file.write(response.content) | |
| print("Model downloaded") | |
| modelPath="./model.gguf" | |
| print(modelPath) | |
| import subprocess | |
| command = ["python3", "-m", "llama_cpp.server", "--model", modelPath, "--host", "0.0.0.0", "--port", "2600", "--n_threads", "4", "--n_gpu_layers","42"] #20 | |
| subprocess.Popen(command) | |
| print("Server ready!") | |
| #url="http://0.0.0.0:2600/v1/completions" | |
| #body={"prompt":"test","max_tokens":1000, "echo":"False","stream":"False"} #e.g. Mixtral-Instruct | |
| #test=requests.post(url, json=body, stream=False) | |
| jina = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-de', trust_remote_code=True, torch_dtype=torch.bfloat16) | |
| #jira.save_pretrained("jinaai_jina-embeddings-v2-base-de") | |
| device='cuda:0' if torch.cuda.is_available() else 'cpu' | |
| jina.to(device) #cuda:0 | |
| print(device) | |
| #----------------- | |
| # ChromaDB-client | |
| #----------------- | |
| class JinaEmbeddingFunction(EmbeddingFunction): | |
| def __call__(self, input: Documents) -> Embeddings: | |
| embeddings = jina.encode(input) #max_length=2048 | |
| return(embeddings.tolist()) | |
| dbPath = "/home/af/Schreibtisch/Code/gradio/Chroma/db/" | |
| onPrem = True if(os.path.exists(dbPath)) else False | |
| if(onPrem==False): dbPath="/home/user/app/db/" | |
| print(dbPath) | |
| client = chromadb.PersistentClient(path=dbPath) | |
| print(client.heartbeat()) | |
| print(client.get_version()) | |
| print(client.list_collections()) | |
| jina_ef=JinaEmbeddingFunction() | |
| embeddingModel=jina_ef | |
| databases=[(date.today(),"0")] # start a list of databases | |
| #--------------------------------------------------------------------- | |
| # Function for formatting single message according to prompt template | |
| #--------------------------------------------------------------------- | |
| def format_prompt0(message, history): | |
| prompt = "<s>" | |
| #for user_prompt, bot_response in history: | |
| # prompt += f"[INST] {user_prompt} [/INST]" | |
| # prompt += f" {bot_response}</s> " | |
| prompt += f"[INST] {message} [/INST]" | |
| return prompt | |
| #------------------------------------------------------------------------- | |
| # Function for formatting multiturn-dialogue according to prompt template | |
| #------------------------------------------------------------------------- | |
| def format_prompt(message, history=None, system=None, RAGAddon=None, system2=None, zeichenlimit=None,historylimit=4, removeHTML=False, | |
| startOfString="<s>", template0=" [INST] {system} [/INST] </s>",template1=" [INST] {message} [/INST]",template2=" {response}</s>"): # mistralai/Mixtral-8x7B-Instruct-v0.1 | |
| #startOfString="<bos>",template0="<start_of_turn>user\n{system}<end_of_turn>\n<start_of_turn>model\n<end_of_turn>\n",template1="<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n",template2="{response}<end_of_turn>\n"): # google/gemma-2-2b-it | |
| #startOfString="<|begin_of_text|><", template0="<|start_header_id|>system<|end_header_id|>\n\nCutting Knowledge Date: December 2023\nToday Date: 26 Jul 2024\n\n{system}\n<|eot_id|>", template1="<|start_header_id|>user<|end_header_id|>\n\n{message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n", template2="{response}</eot_id>"): # meta-llama/Llama-3.1-8B-Instruct | |
| if zeichenlimit is None: zeichenlimit=1000000000 # :-) | |
| prompt = "" | |
| if RAGAddon is not None: | |
| system += RAGAddon | |
| if system is not None: | |
| prompt += template0.format(system=system) #"<s>" | |
| if history is not None: | |
| for user_message, bot_response in history[-historylimit:]: | |
| if user_message is None: user_message = "" | |
| if bot_response is None: bot_response = "" | |
| bot_response = re.sub("\n\n<details>((.|\n)*?)</details>","", bot_response) # remove RAG-compontents | |
| if removeHTML==True: bot_response = re.sub("<(.*?)>","\n", bot_response) # remove HTML-components in general (may cause bugs with markdown-rendering) | |
| if user_message is not None: prompt += template1.format(message=user_message[:zeichenlimit]) | |
| if bot_response is not None: prompt += template2.format(response=bot_response[:zeichenlimit]) | |
| if message is not None: prompt += template1.format(message=message[:zeichenlimit]) | |
| if system2 is not None: | |
| prompt += system2 | |
| return startOfString+prompt | |
| #-------------------------------------------- | |
| # Function for converting pdf-files to text | |
| #-------------------------------------------- | |
| def convertPDF(pdf_file, allow_ocr=False): | |
| reader = PdfReader(pdf_file) | |
| full_text = "" | |
| page_list = [] | |
| def extract_text_from_pdf(reader): | |
| full_text = "" | |
| page_list = [] | |
| page_count = 1 | |
| for idx, page in enumerate(reader.pages): | |
| text = page.extract_text() | |
| if len(text) > 0: | |
| page_list.append(text) | |
| #full_text += f"---- Page {idx} ----\n" + text + "\n\n" | |
| page_count += 1 | |
| return full_text.strip(), page_count, page_list | |
| # Check if there are any images | |
| image_count = sum(len(page.images) for page in reader.pages) | |
| # If there are images and not much content, you may want to perform OCR on the document | |
| if allow_ocr: | |
| print(f"{image_count} Images") | |
| if image_count > 0 and len(full_text) < 1000: | |
| out_pdf_file = pdf_file.replace(".pdf", "_ocr.pdf") | |
| ocrmypdf.ocr(pdf_file, out_pdf_file, force_ocr=True) | |
| reader = PdfReader(out_pdf_file) | |
| # Extract text: | |
| full_text, page_count, page_list = extract_text_from_pdf(reader) | |
| l = len(page_list) | |
| print(f"{l} Pages") | |
| # Extract metadata | |
| metadata = { | |
| "author": reader.metadata.author, | |
| "creator": reader.metadata.creator, | |
| "producer": reader.metadata.producer, | |
| "subject": reader.metadata.subject, | |
| "title": reader.metadata.title, | |
| "image_count": image_count, | |
| "page_count": page_count, | |
| "char_count": len(full_text), | |
| } | |
| return page_list, full_text, metadata | |
| #------------------------------------------ | |
| # Function for splitting text with overlap | |
| #------------------------------------------ | |
| def split_with_overlap0(text,chunk_size=3500, overlap=700): | |
| """ Split text in chunks based on number of characters (chunk_size) with chunks overlapping (overlap)""" | |
| chunks=[] | |
| step=max(1,chunk_size-overlap) | |
| for i in range(0,len(text),step): | |
| end=min(i+chunk_size,len(text)) | |
| chunks.append(text[i:end]) | |
| return chunks | |
| import re | |
| def split_with_overlap(text, chunk_size=3500, overlap=700, pattern=r'([.!;?][ \n\r]|[\n\r]{2,})', variant=1, verbose=False): | |
| """ Split text in chunks based on regex (pattern) matches. By default the pattern is '([.!;?][ \\n\\r]|[\\n\\r]{2,})' Chunks are no longer than a certain number of characters (chunk_size) with chunks overlapping (overlap). | |
| By default (variant=1) chunking is based on complete sentences, but it's also possible to split only within the left overlap region and within the rest of the chunk-size (variant==2) or strictly within both overlap-regions (variant=3). | |
| """ | |
| chunks = [] | |
| overlap=min(overlap,chunk_size) # Overlap kann nicht größer sein als chunk_size | |
| step = max(1, chunk_size - overlap) # step richtet sich nach chunk_size und overlap | |
| def find_pattern(text): # Funktion zur Suche nach dem Muster | |
| return re.search(pattern, text) | |
| i, lastEnd = 0,0 | |
| while i<len(text): | |
| print("i="+str(i)) | |
| end = min(i + chunk_size, len(text)) | |
| pattern_match = find_pattern(text[i:end]) # erstes Vorkommnis (if any) | |
| matchesStart = [x.start() for x in re.finditer(pattern, text[i:end])] # start aller matches | |
| matchesEnd = [x.start() for x in re.finditer(pattern, text[i:end])] # end aller matches | |
| step = max(1, chunk_size - overlap) # Normalerweise beträgt ein Step chunk_size - overlap | |
| if pattern_match: # Wenn (mindestens) ein Satzzeichen gefunden wurde | |
| for s in matchesStart: # gehe jedes Satzzeichen durch | |
| if ((variant<=2 and s>=overlap) or (variant==3 and s>=overlap and s>(chunk_size-overlap))): # wenn das Satzzeichen nicht im Overlap links liegt (1) oder zusätzlich im reechten Overlap liegt (2) - wobei letzteres unvollständige Sätze bedeuten kann | |
| end=s+i+1 # Setze end auf den Start des Patterns/Satzzeichens im gesamten Text | |
| if(verbose==True): print("***move end:"+str(end)+"; step="+str(step)) | |
| if(s<(chunk_size-overlap)):step=min(step,max(1,s-overlap)) # Springe mit step höchstens zum Ende des Satzzeichens (nur erforderlich, wenn end nicht im Overlap) | |
| if ((variant==1 and i>0) or (variant>=2 and pattern_match.start()<overlap and i>0)): # wenn das erste Satzzeichen im Overlap liegt | |
| i=i+pattern_match.start()+1 # Verzichte auf Textteile vor dem ersten Satzzeichen | |
| if(verbose==True): print("i="+str(i)+"; end="+str(end)+"; step="+str(step)+"; len="+str(len(text))+"; match="+str(pattern_match)+"; text="+text[i:end]+"; rest="+text[end:]) | |
| if(end>lastEnd): # wenn das Ende sich verschoben hat (und nicht nur den Satzbeginn zu einem bereits bekannten Satz abschneidet) | |
| chunks.append(text[i:end]) | |
| lastEnd=end | |
| if(verbose==True): print("Text at position "+str(i)+": "+text[i:end]) | |
| i += step | |
| if(len(text[end:])>0): chunks.append(text[end:]) # Ergänze am ende etwaigen Rest | |
| return chunks | |
| fiveChars= "(?<![ \n\(]bspw|[ \n]inkl)" | |
| fourChars= "(?<![ \n\(]sog|[ \n]Mio|[ \n]Mrd|[ \n]Tsd|[ \n]Tel)" | |
| threeChars= "(?<!www|bzw|etc|ggf|[ \n\(]al|[ \n\(]St|[ \n\(]dh|[ \n\(]va|[ \n\(]ca|[ \n\(]Dr|[ \n\(]Hr|[ \n\(]Fr|[0-9]ff)" | |
| twoChars= "(?<![ \n\(][A-Za-zΆ-Ωά-ωäöüß])" | |
| oneChars= "(?<![0-9.])" | |
| sentenceRegex="(?<=[^.]{4})"+fiveChars+fourChars+threeChars+twoChars+oneChars+"[.?!](?![A-Za-zΆ-Ωά-ωäöüß0-9.!?'\"])" | |
| sectionRegex="\n[ ]*\n[\n ]*" | |
| splitRegex="("+sentenceRegex+"|"+sectionRegex+")" | |
| #--------------------------------------------------------------- | |
| # Function for adding docs to ChromaDB and/or return collection | |
| #--------------------------------------------------------------- | |
| def add_doc(path, session): | |
| global device | |
| print("def add_doc!") | |
| print(path) | |
| anhang=False | |
| if(str.lower(path).endswith(".pdf") and os.path.exists(path)): | |
| doc=convertPDF(path) | |
| if(len(doc[0])>5): | |
| if(not "cuda" in device): | |
| doc="\n\n".join(doc[0][0:5]) | |
| gr.Info("PDF uploaded to DB_"+str(session)+", start Indexing excerpt (demo-mode: first 5 pages on CPU setups)!") | |
| else: | |
| doc="\n\n".join(doc[0]) | |
| gr.Info("PDF uploaded to DB_"+str(session)+", start Indexing!") | |
| else: | |
| doc="\n\n".join(doc[0]) | |
| gr.Info("PDF uploaded to DB_"+str(session)+", start Indexing!") | |
| anhang=True | |
| else: | |
| gr.Info("No PDF attached - answer based on DB_"+str(session)+".") | |
| client = chromadb.PersistentClient(path=dbPath) | |
| print(str(client.list_collections())) | |
| print(str(session)) | |
| dbName="DB_"+str(session) | |
| if(not "name="+dbName in str(client.list_collections())): | |
| # client.delete_collection(name=dbName) | |
| collection = client.create_collection( | |
| name=dbName, | |
| embedding_function=embeddingModel, | |
| metadata={"hnsw:space": "cosine"}) | |
| else: | |
| collection = client.get_collection( | |
| name=dbName, embedding_function=embeddingModel) | |
| if(anhang==True): | |
| corpus=split_with_overlap(doc,3500,700,pattern=splitRegex) | |
| print("Length of corpus: "+str(len(corpus))) | |
| print("Corpus:"+str(corpus)) | |
| then = datetime.now() | |
| x=collection.get(include=[])["ids"] | |
| print(len(x)) | |
| if(len(x)==0): | |
| chunkSize=40000 | |
| for i in range(round(len(corpus)/chunkSize+0.5)): #0 is first batch, 3 is last (incomplete) batch given 133497 texts | |
| print("embed batch "+str(i)+" of "+str(round(len(corpus)/chunkSize+0.5))) | |
| ids=list(range(i*chunkSize,(i*chunkSize+chunkSize))) | |
| batch=corpus[i*chunkSize:(i*chunkSize+chunkSize)] | |
| textIDs=[str(id) for id in ids[0:len(batch)]] | |
| ids=[str(id+len(x)+1) for id in ids[0:len(batch)]] # id refers to chromadb-unique ID | |
| collection.add(documents=batch, ids=ids, | |
| metadatas=[{"date": str("2024-10-10")} for b in batch]) #"textID":textIDs, "id":ids, | |
| print("finished batch "+str(i)+" of "+str(round(len(corpus)/40000+0.5))) | |
| now = datetime.now() | |
| gr.Info(f"Indexing complete!") | |
| print(now-then) #zu viel GB für sentences (GPU), bzw. 0:00:10.375087 für chunks | |
| return(collection) | |
| #-------------------------------------------------------- | |
| # Function for response to user queries and pot. addenda | |
| #-------------------------------------------------------- | |
| def multimodal_response(message, history, dropdown, hfToken, request: gr.Request): | |
| print("def multimodal response!") | |
| if(hfToken.startswith("hf_")): # use HF-hub with custom token if token is provided | |
| inferenceClient = InferenceClient(model=myModel, token=hfToken) | |
| else: | |
| inferenceClient = InferenceClient(myModel) | |
| global databases | |
| if request: | |
| session=request.session_hash | |
| else: | |
| session="0" | |
| length=str(len(history)) | |
| print(databases) | |
| if(not databases[-1][1]==session): | |
| databases.append((date.today(),session)) | |
| #print(databases) | |
| query=message["text"] | |
| if(len(message["files"])>0): # is there at least one file attached? | |
| collection=add_doc(message["files"][0], session) | |
| else: # otherwise, you still want to get the collection with the session-based db | |
| collection=add_doc(message["text"], session) | |
| client = chromadb.PersistentClient(path=dbPath) | |
| print(str(client.list_collections())) | |
| x=collection.get(include=[])["ids"] | |
| ragQuery=[format_prompt(query, history, historylimit=2, | |
| #startOfString="", template0="{system}\n",template1="USER: {message}\n\n",template2="ASSISTANT: {response}\n\n") if len(history)>0 else query] # embed simply-formated dialogue | |
| startOfString="", template1="{message}\n\n",template2="") if len(history)>0 else query] # embed simple string of User-queries only | |
| context=collection.query(query_texts=ragQuery, n_results=3) | |
| #context=["<Kontext "+str(i)+"> "+str(c)+"</Kontext "+str(i)+">" for i,c in enumerate(context["documents"][0])] | |
| context=["Kontext "+str(i+1)+": \""+re.sub("\"","'",str(c))+"\"" for i,c in enumerate(context["documents"][0])] | |
| gr.Info("Kontext:\n"+str(context)) | |
| generate_kwargs = dict( | |
| temperature=float(0.9), | |
| max_new_tokens=5000, | |
| top_p=0.95, | |
| repetition_penalty=1.0, | |
| do_sample=True, | |
| seed=42, | |
| ) | |
| system="Mit Blick auf das folgende Gespräch und den relevanten Kontext, antworte auf die aktuelle Frage des Nutzers. "+\ | |
| "Antworte ausschließlich auf Basis der Informationen im Kontext.\n\nKontext:\n\n"+\ | |
| str("\n\n".join(context)) | |
| #"Given the following conversation, relevant context, and a follow up question, "+\ | |
| #"reply with an answer to the current question the user is asking. "+\ | |
| #"Return only your response to the question given the above information "+\ | |
| #"following the users instructions as needed.\n\nContext:"+\ | |
| print(system) | |
| #formatted_prompt = format_prompt0(system+"\n"+query, history) | |
| formatted_prompt = format_prompt(query, history,system=system) | |
| print(formatted_prompt) | |
| output = "" | |
| if(not "GGUF" in myModel): | |
| try: | |
| stream = inferenceClient.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False) | |
| for response in stream: | |
| output += response.token.text | |
| yield output | |
| except Exception as e: | |
| output = "Für weitere Antworten von der KI gebe bitte einen gültigen HuggingFace-Token an." | |
| if(len(context)>0): | |
| output += "\nBis dahin helfen dir hoffentlich die folgenden Quellen weiter:" | |
| yield output | |
| print(str(e)) | |
| else: | |
| try: | |
| #generate_kwargs["prompt"]=formatted_prompt # | |
| generate_kwargs={"prompt":formatted_prompt,"max_tokens":1000, "echo":"False","stream":"True"} #e.g. Mixtral-Instruct | |
| url="http://0.0.0.0:2600/v1/completions" | |
| response="" | |
| buffer="" | |
| print("URL: "+url) | |
| print("User: "+str(message)+"\nAssistant: ") | |
| for text in requests.post(url, json=generate_kwargs, stream=True): #-H 'accept: application/json' -H 'Content-Type: application/json' | |
| if buffer is None: buffer="" | |
| buffer=str("".join(buffer)) | |
| text=text.decode('utf-8') | |
| if((text.startswith(": ping -")==False) & (len(text.strip("\n\r"))>0)): buffer=buffer+str(text) | |
| buffer=buffer.split('"finish_reason": null}]}') | |
| if(len(buffer)==1): | |
| buffer="".join(buffer) | |
| pass | |
| if(len(buffer)==2): | |
| part=buffer[0]+'"finish_reason": null}]}' | |
| if(part.lstrip('\n\r').startswith("data: ")): part=part.lstrip('\n\r').replace("data: ", "") | |
| try: | |
| part = str(json.loads(part)["choices"][0]["text"]) | |
| print(part, end="", flush=True) | |
| output += part | |
| buffer="" | |
| except Exception as e: | |
| print("Exception:"+str(e)) | |
| pass | |
| yield output | |
| except Exception as e: | |
| output = "Die KI antwortet gerade nicht." | |
| if(len(context)>0): | |
| output += "\nBis dahin helfen dir hoffentlich die folgenden Quellen weiter:" | |
| yield output | |
| print(str(e)) | |
| if(len(context)>0): | |
| output=output+"\n\n<br><details open><summary><strong>Quellen</strong></summary><br><ul>"+ "".join(["<li>" + c + "</li>" for c in context])+"</ul></details>" | |
| yield output | |
| #------------------------------ | |
| # Launch Gradio-ChatInterface | |
| #------------------------------ | |
| i=gr.ChatInterface(multimodal_response, | |
| title="Frag dein PDF", | |
| multimodal=True, | |
| additional_inputs=[ | |
| gr.Dropdown( | |
| info="Wähle eine Variante", | |
| choices=["1","2","3"], | |
| value="1", | |
| label="Variante"), | |
| gr.Textbox( | |
| value="", | |
| label="HF_token"), | |
| ]) | |
| i.launch() #allowed_paths=["."]) | |