| import gradio as gr | |
| from datasets import load_dataset, load_from_disk, Dataset | |
| from transformers import AutoTokenizer, AutoModel | |
| import torch | |
| import pandas as pd | |
| import base64 | |
| import html | |
| import re | |
| from functools import * | |
| model_ckpt = "nomic-ai/nomic-embed-text-v1.5" | |
| model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1" | |
| model_ckpt = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1" | |
| tokenizer = AutoTokenizer.from_pretrained(model_ckpt) | |
| model = AutoModel.from_pretrained(model_ckpt, trust_remote_code=True) | |
| device = torch.device("cpu") | |
| model.to(device) | |
| def cls_pooling(model_output): | |
| return model_output.last_hidden_state[:, 0] | |
| def get_embeddings(text_list): | |
| encoded_input = tokenizer( | |
| text_list, padding=True, truncation=True, return_tensors="pt" | |
| ) | |
| encoded_input = {k: v.to(device) for k, v in encoded_input.items()} | |
| model_output = model(**encoded_input) | |
| return cls_pooling(model_output) | |
| embeddings_dataset = load_dataset("Vadim212/doctest2")["train"] | |
| embeddings_dataset.add_faiss_index(column="embeddings") | |
| def filter_f(x, field, value): | |
| return x[field] == value or x[field] is None | |
| def filter_product(question, dataset): | |
| product_line = "" | |
| for pl in ["report", "dashboard", "form"]: | |
| if re.search(f"\\b({pl}|{pl}s|{pl}`s)\\b", question, re.IGNORECASE) is not None: product_line = f"{pl}s" | |
| product = "" | |
| if re.search("\\b(javascript|java script)\\b", question, re.IGNORECASE) is not None: product = "js" | |
| for pr in ["web", "net", "js", "wpf", "php", "blazor", "angular", "java"]: | |
| if re.search(f"{pr}\\b", question, re.IGNORECASE) is not None: product = pr | |
| if (product_line != "") or (product != ""): | |
| product_line_filter = partial(filter_f, field = "product_line", value = product_line) if product_line != "" else lambda x: True | |
| product_filter = partial(filter_f, field = "product", value = product) if product != "" else lambda x: True | |
| if len(dataset.list_indexes()) > 0 : dataset.drop_index("embeddings") | |
| dataset = dataset.filter(lambda x: product_line_filter(x) and product_filter(x)) | |
| return dataset | |
| def get_html(row): | |
| pr = row["product"] | |
| product_ = f" [{pr}]" if pr is not None else "" | |
| product_line = f" [{row.product_line}]" if row.product_line is not None else "" | |
| path = f"[{row.path}]" if row.path is not None and len(row.path) > 0 else "" | |
| result = f"""<div style='font-size:16pt'>{html.escape(row.title)} <font style='font-size:10pt'>[{row.origin}]{product_line}{product_}({row.scores})</font> </div> | |
| <div style='font-size:8pt'>{html.escape(path)}</div> | |
| {html.escape(row.content)}<br><a href='{row.url}' target='_blank'>Link</a>""" | |
| return result | |
| def find(question): | |
| question_embedding = get_embeddings([question]).cpu().detach().numpy() | |
| dataset = filter_product(question, embeddings_dataset) | |
| if len(dataset.list_indexes()) == 0: dataset.add_faiss_index(column="embeddings") | |
| scores, samples = dataset.get_nearest_examples( | |
| "embeddings", question_embedding, k=20 | |
| ) | |
| samples_df = pd.DataFrame.from_dict(samples) | |
| samples_df["scores"] = scores | |
| samples_df.sort_values("scores", ascending=True, inplace=True) | |
| samples_df.drop_duplicates(subset=["url"], inplace=True) | |
| samples_df = samples_df[:10] | |
| result = [get_html(row) for i, row in samples_df.iterrows()] | |
| return result | |
| demo = gr.Blocks() | |
| with demo: | |
| inp = gr.Textbox(placeholder="Enter prompt",label= "Prompt like: 'how to export to PDF?', 'What is report?', 'blazor designer', 'how to render report?', 'List of export formats', 'Supported databases'" ) | |
| find_btn = gr.Button("Find") | |
| big_block = [gr.HTML("") for i in range(10)] | |
| find_btn.click(find, | |
| inputs=inp, | |
| outputs=big_block) | |
| demo.launch(inline=True, width=400) |