georeactor/asknyc-vectorsearch

1import os2 3import cohere4import gradio as gr5import numpy as np6import pinecone7import torch8from transformers import AutoModel, AutoTokenizer9 10co = cohere.Client(os.environ.get('COHERE_API', ''))11pinecone.init(12    api_key=os.environ.get('PINECONE_API', ''),13    environment=os.environ.get('PINECONE_ENV', '')14)15 16# model = AutoModel.from_pretrained('monsoon-nlp/gpt-nyc')17# tokenizer = AutoTokenizer.from_pretrained('monsoon-nlp/gpt-nyc')18# zos = np.zeros(4096-1024).tolist()19 20def list_me(matches):21    result = ''22    for match in matches:23        result += '<li><a target="_blank" href="https://reddit.com/r/AskNYC/comments/' + match['id'] + '">'24        result += match['metadata']['question']25        result += '</a>'26        if 'body' in match['metadata']:27            result += '<br/>' + match['metadata']['body']28        result += '</li>'29    return result.replace('/mini', '/')30 31 32def query(question):33    # Cohere search34    response = co.embed(35        model='large',36        texts=[question],37    )38    index = pinecone.Index("gptnyc")39    closest = index.query(40        top_k=2,41        include_metadata=True,42        vector=response.embeddings[0],43    )44 45    # SGPT search46    # batch_tokens = tokenizer(47    #     [question],48    #     padding=True,49    #     truncation=True,50    #     return_tensors="pt"51    # )52    # with torch.no_grad():53    #     last_hidden_state = model(**batch_tokens, output_hidden_states=True, return_dict=True).last_hidden_state54    # weights = (55    #     torch.arange(start=1, end=last_hidden_state.shape[1] + 1)56    #     .unsqueeze(0)57    #     .unsqueeze(-1)58    #     .expand(last_hidden_state.size())59    #     .float().to(last_hidden_state.device)60    # )61    # input_mask_expanded = (62    #     batch_tokens["attention_mask"]63    #     .unsqueeze(-1)64    #     .expand(last_hidden_state.size())65    #     .float()66    # )67    # sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded * weights, dim=1)68    # sum_mask = torch.sum(input_mask_expanded * weights, dim=1)69    # embeddings = sum_embeddings / sum_mask70    # closest_sgpt = index.query(71    #     top_k=2,72    #     include_metadata=True,73    #     namespace="mini",74    #     vector=embeddings[0].tolist() + zos,75    # )76 77    return '<h3>Cohere</h3><ul>' + list_me(closest['matches']) + '</ul>'78    #'<h3>SGPT</h3><ul>' + list_me(closest_sgpt['matches']) + '</ul>'79 80 81iface = gr.Interface(82    fn=query,83    inputs="text",84    outputs="html"85)86iface.launch()87