query sgpt

app.py

@@ -2,7 +2,10 @@ import os
 
 import cohere
 import gradio as gr
+import numpy as np
 import pinecone
+import torch
+from transformers import AutoModel, AutoTokenizer
 
 co = cohere.Client(os.environ.get('COHERE_API', ''))
 pinecone.init(
@@ -10,6 +13,10 @@ pinecone.init(
     environment=os.environ.get('PINECONE_ENV', '')
 )
 
+model = AutoModel.from_pretrained('monsoon-nlp/gpt-nyc')
+tokenizer = AutoTokenizer.from_pretrained('monsoon-nlp/gpt-nyc')
+zos = np.zeros(4096-1024).tolist()
+
 def list_me(matches):
     result = ''
     for match in matches:
@@ -19,10 +26,11 @@ def list_me(matches):
         if 'body' in match['metadata']:
             result += '<br/>' + match['metadata']['body']
         result += '</li>'
-    return result
+    return result.replace('/mini', '/')
 
 
 def query(question):
+    # Cohere search
     response = co.embed(
         model='large',
         texts=[question],
@@ -34,7 +42,39 @@ def query(question):
         vector=response.embeddings[0],
     )
 
-    return '<ul>' + list_me(closest['matches']) + '</ul>'
+    # SGPT search
+    batch_tokens = tokenizer(
+        [question],
+        padding=True,
+        truncation=True,
+        return_tensors="pt"
+    )
+    with torch.no_grad():
+        last_hidden_state = model(**batch_tokens, output_hidden_states=True, return_dict=True).last_hidden_state
+    weights = (
+        torch.arange(start=1, end=last_hidden_state.shape[1] + 1)
+        .unsqueeze(0)
+        .unsqueeze(-1)
+        .expand(last_hidden_state.size())
+        .float().to(last_hidden_state.device)
+    )
+    input_mask_expanded = (
+        batch_tokens["attention_mask"]
+        .unsqueeze(-1)
+        .expand(last_hidden_state.size())
+        .float()
+    )
+    sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded * weights, dim=1)
+    sum_mask = torch.sum(input_mask_expanded * weights, dim=1)
+    embeddings = sum_embeddings / sum_mask
+    closest_sgpt = index.query(
+        top_k=2,
+        include_metadata=True,
+        namespace="mini",
+        vector=embeddings[0].tolist() + zos,
+    )
+
+    return '<h3>Cohere</h3><ul>' + list_me(closest['matches']) + '</ul><h3>SGPT</h3><ul>' + list_me(closest_sgpt['matches']) + '</ul>'
 
 
 iface = gr.Interface(

requirements.txt

@@ -1,2 +1,4 @@
 cohere==3.10.0
 pinecone-client==2.2.1
+torch
+transformers==4.26.1