|
import gradio as gr |
|
import spaces |
|
import torch |
|
from datasets import load_dataset |
|
from sentence_transformers import SentenceTransformer |
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
import os |
|
|
|
dataset = load_dataset("ariG23498/pis-blogs-chunked") |
|
embedding_model = SentenceTransformer(model_name_or_path="all-mpnet-base-v2") |
|
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it", key="HF_TOKEN", value=os.getenv("auth")) |
|
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.bfloat16, key="HF_TOKEN", value=os.getenv("auth")) |
|
|
|
@spaces.GPU |
|
def process_query(query): |
|
embedding_model = embedding_model.to("cuda") |
|
text_embeddings = embedding_model.encode(dataset["train"]["text"]) |
|
|
|
query_embedding = embedding_model.encode(query) |
|
similarity_scores = embedding_model.similarity(query_embedding, text_embeddings) |
|
top_indices = (-similarity_scores).argsort()[0][:5] |
|
|
|
context = dataset["train"]["text"][top_indices[0]] |
|
url = dataset["train"]["url"][top_indices[0]] |
|
|
|
input_text = ( |
|
f"Based on the context provided, '{context}', how would" |
|
f"you address the user's query regarding '{query}'? Please" |
|
" provide a detailed and contextually relevant response." |
|
) |
|
|
|
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda") |
|
len_text = len(input_text) |
|
model = model.to("cuda") |
|
with torch.inference_mode(): |
|
generated_outputs = model.generate(**input_ids, max_new_tokens=1000, do_sample=False) |
|
generated_outputs = tokenizer.batch_decode(generated_outputs, skip_special_tokens=True) |
|
|
|
response = generated_outputs[0][len_text:] |
|
return url, response |
|
|
|
demo = gr.Interface( |
|
fn=process_query, |
|
inputs=gr.Textbox(label="User Query"), |
|
outputs=[gr.Textbox(label="URL"), gr.Textbox(label="Generated Response")] |
|
) |
|
|
|
demo.launch() |
|
|
