File size: 1,820 Bytes
33ffdb4 a3f8683 33ffdb4 20309d7 ee678e0 20309d7 bd41503 33ffdb4 ee678e0 33ffdb4 33a95fe 33ffdb4 |
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 |
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
os.environ["HF_TOKEN"] = os.getenv("auth")
dataset = load_dataset("ariG23498/pis-blogs-chunked")
embedding_model = SentenceTransformer(model_name_or_path="all-mpnet-base-v2", device="cuda")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.bfloat16, device_map="auto")
@spaces.GPU(duration=300)
def process_query(query):
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)
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()
|