chore: initial commit

app.py

@@ -0,0 +1,51 @@
+import gradio as gr
+import spaces
+import torch
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+@spaces.GPU
+def process_query(query):
+    dataset = load_dataset("ariG23498/pis-blogs-chunked")
+    embedding_model = SentenceTransformer(model_name_or_path="all-mpnet-base-v2", device="cuda:0")
+    text_embeddings = embedding_model.encode(dataset["train"]["text"])
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
+    model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.bfloat16, device_map="auto")
+
+    print(f"Query: {query}")
+    
+    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]]
+    
+    print(f"Searching URL: {url}")
+    print(f"Found context: {context}")
+    
+    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:0")
+    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()

requirements.txt

@@ -0,0 +1,5 @@
+accelerate
+transformers
+gradio
+sentence-transformers
+datasets