US_History_QA / app.py
RyanDA's picture
Update app.py
352215e
import gradio as gr
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain.vectorstores import Chroma
import openai
import os
openai.api_key = os.environ["OPENAI_API_KEY"]
embedding = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
vectordb = Chroma(persist_directory='db',
embedding_function=embedding)
retriever = vectordb.as_retriever(search_kwargs={"k": 4})
def answer(message):
ctx = retriever.get_relevant_documents(message)
prompt = """Use the following pieces of context to answer the question at the end.
If you don't know the answer, just say that you don't know, don't try to
make up an answer.
{context}
Question: {question}
Helpful Answer: """.format(context = " ".join([doc.page_content for doc in ctx]), question = message)
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo",
messages= [
{'role': 'user', 'content': prompt}
],
temperature=0
)
sources = "\n\nSources:\n" + "\n".join([doc.metadata['source'] for doc in ctx])
ans = response['choices'][0]['message']['content'] + sources
return ans
iface = gr.Interface(fn=answer, inputs="text", outputs="text", examples=["When was the declaration of independence signed?", "What year did the United States enter World War I?", "What was considered the peak of tension within the Cold War?"])
iface.launch()