mirror of
https://github.com/wassname/stampy-chat.git
synced 2026-09-20 13:20:52 +08:00
123 lines
4.7 KiB
Python
123 lines
4.7 KiB
Python
import json
|
|
import pickle
|
|
from typing import List, Tuple
|
|
import os
|
|
|
|
|
|
import numpy as np # TODO: Add to requirements.txt
|
|
from tenacity import ( # TODO: Add to requirements.txt
|
|
retry,
|
|
stop_after_attempt,
|
|
wait_random_exponential,
|
|
)
|
|
import openai # TODO: Add to requirements.txt
|
|
|
|
|
|
os.environ.get('OPENAI_API_KEY')
|
|
openai.api_key = os.environ.get('OPENAI_API_KEY')
|
|
|
|
from pathlib import Path # BAD
|
|
|
|
EMBEDDING_MODEL = "text-embedding-ada-002" # BAD
|
|
COMPLETIONS_MODEL = "text-davinci-003" # BAD
|
|
|
|
LEN_EMBEDDINGS = 1536 # BAD
|
|
MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. # BAD
|
|
|
|
project_path = Path(__file__).parent.parent.parent
|
|
PATH_TO_DATA = project_path / "src" / "data" / "alignment_texts.jsonl" # Path to the dataset .jsonl file. # BAD
|
|
PATH_TO_EMBEDDINGS = project_path / "src" / "data" / "embeddings.npy" # Path to the saved embeddings (.npy) file. # BAD
|
|
PATH_TO_DATASET = project_path / "src" / "data" / "dataset.pkl" # Path to the saved dataset (.pkl) file, containing the dataset class object. # BAD
|
|
|
|
from http.server import BaseHTTPRequestHandler
|
|
|
|
class handler(BaseHTTPRequestHandler):
|
|
|
|
# post request = calculate factorial of passed number
|
|
def do_POST(self):
|
|
self.send_response(200)
|
|
self.send_header('Content-type', 'application/json')
|
|
self.end_headers()
|
|
content_length = int(self.headers['Content-Length'])
|
|
post_data = self.rfile.read(content_length)
|
|
data = json.loads(post_data)
|
|
|
|
results = {};
|
|
|
|
for i, link in enumerate(embeddings(data['query'])):
|
|
results[i] = json.dumps(link.__dict__)
|
|
|
|
self.wfile.write(json.dumps(results).encode('utf-8'))
|
|
|
|
|
|
class Link:
|
|
def __init__(self, url, title):
|
|
self.url = url
|
|
self.title = title
|
|
|
|
|
|
@retry(wait=wait_random_exponential(min=1, max=10), stop=stop_after_attempt(4))
|
|
def get_embedding(text: str) -> np.ndarray:
|
|
result = openai.Embedding.create(model=EMBEDDING_MODEL, input=text)
|
|
return result["data"][0]["embedding"]
|
|
|
|
def get_top_k_blocks(user_query: str, k: int, HyDE: bool = False) -> List[str]:
|
|
"""Get the top k blocks that are most semantically similar to the query, using the provided dataset.
|
|
|
|
Args:
|
|
query (str): The query to be searched for.
|
|
k (int): The number of blocks to return.
|
|
HyDE (bool, optional): Whether to use HyDE or not. Defaults to False.
|
|
|
|
Returns:
|
|
List[str]: A list of the top k blocks that are most semantically similar to the query.
|
|
"""
|
|
# Get the dataset
|
|
with open(PATH_TO_DATASET, "rb") as f:
|
|
dataset = pickle.load(f)
|
|
|
|
# Get the embedding for the query.
|
|
query_embedding = get_embedding(user_query)
|
|
|
|
# If HyDE is enabled, produce a no-context ChatCompletion to the query.
|
|
if HyDE:
|
|
messages = [
|
|
{"role": "system", "content": "You are a knowledgeable AI Alignment assistant. Do your best to answer the user's question, even if you don't know the answer for sure."},
|
|
{"role": "user", "content": user_query},
|
|
]
|
|
HyDE_completion = openai.ChatCompletion.create(
|
|
model=COMPLETIONS_MODEL,
|
|
messages=messages,
|
|
temperature=0.0,
|
|
max_tokens=200
|
|
)["choices"][0]["text"]
|
|
HyDe_completion_embedding = get_embedding(f"Question: {user_query}\n\nAnswer: {HyDE_completion}")
|
|
|
|
similarity_scores = np.dot(dataset.metadataset.embeddings, HyDe_completion_embedding)
|
|
else:
|
|
similarity_scores = np.dot(dataset.metadataset.embeddings, query_embedding)
|
|
|
|
ordered_blocks = np.argsort(similarity_scores)[::-1] # Sort the blocks by similarity score
|
|
top_k_indices = ordered_blocks[:k] # Get the top k indices
|
|
top_k = [dataset.metadataset.embedding_strings[i] for i in top_k_indices] # Get the top k strings
|
|
|
|
# Get associated links
|
|
|
|
return top_k
|
|
|
|
def embeddings(query):
|
|
# write a function here that takes a query, returns a bunch of semantically similar links
|
|
|
|
|
|
return [ \
|
|
Link('https://www.lesswrong.com/posts/FinfRNLMfbq5ESxB9/microsoft-research-paper-claims-sparks-of-artificial', \
|
|
'Microsoft Research Paper Claims Sparks of Artificial Intelligence'), \
|
|
Link('https://www.lesswrong.com/posts/XhfBRM7oRcpNZwjm8/abstracts-should-be-either-actually-short-tm-or-broken-into', \
|
|
'Abstracts should be either actually short™ or broken into'), \
|
|
Link('https://www.lesswrong.com/posts/ohXcBjGvazPAxq2ex/continue-working-on-hard-alignment-don-t-give-up', \
|
|
'Continue working on hard alignment, don\'t give up'), \
|
|
Link('https://www.lesswrong.com/posts/' + query + '/this-is-a-test', \
|
|
'This is a test of ' + query), \
|
|
]
|
|
|