InsightAgent / utils.py
lichengen's picture
first commit
8d1c3e1
Raw
History Blame Contribute Delete
2.31 kB
import openai, tiktoken, os
import pandas as pd
from openai.embeddings_utils import get_embedding, cosine_similarity
# embedding model parameters
embedding_model = "text-embedding-ada-002"
embedding_encoding = "cl100k_base" # this the encoding for text-embedding-ada-002
max_tokens = 8000 # the maximum for text-embedding-ada-002 is 8191
encoding = tiktoken.get_encoding("cl100k_base")
def set_openai_api_key(api_key):
if api_key and api_key.startswith("sk-") and len(api_key) > 50:
openai.api_key = api_key
else:
raise gr.Error("OpenAI API key incorrect.")
# Prepare prompt
def prepare_prompt(prompt, results):
tokens_limit = 16000 # Limit for gpt-3.5-turbo-16k
user_start = (
"請只根據下述的文本,使用繁體中文(zh-TW)回答問題。排除相似重複的語意,用字精鍊而清晰,出現過一次的人名就不需要一直顯示全名以及完整職稱。請嚴格遵守只根據下述的文本的規定,如果詢問的問題超過文本的範圍,請回答你不知道。\n\n"+
"文本:\n"
)
user_end = (
f"\n\n問題: {prompt}\n 答案:"
)
system = """
你是一個萬能文字助手,你擅長從大量的文章中,辨識出相關主題,並整理成重點摘要。
"""
count_of_tokens_consumed = len(encoding.encode("\"role\":\"system\"" +
"\"content\" :\"" + system +
user_start + "\n\n---\n\n" + user_end ))
count_of_tokens_for_context = tokens_limit - count_of_tokens_consumed
contexts =""
# Fill in context as long as within limit
for i in range(len(results)):
if (count_of_tokens_for_context>=results.n_tokens.iloc[i]):
contexts += results.text.iloc[i] + "\n"
count_of_tokens_for_context -=1
count_of_tokens_for_context -= results.n_tokens.iloc[i]
complete_prompt = user_start + contexts + "\n\n---\n\n" + user_end
return complete_prompt
def answer(messages):
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo-16k",
messages=messages,
temperature=0,
stream=True
)
return response