You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
94 lines
3.7 KiB
94 lines
3.7 KiB
from _chroma import ChromaDB |
|
from _openai import LLM_OpenAI as LLM |
|
import streamlit as st |
|
from print_color import * |
|
|
|
def get_docs(user_input): |
|
docs = chroma.query('mala_interrogations', user_input, n_results=5) |
|
return docs |
|
|
|
def generate_prompt(user_input, docs): |
|
texts = [text for text in docs['documents'][0]] |
|
metas = [{'person': meta['person'], 'date': meta['date']} for meta in docs['metadatas'][0]] |
|
combined_data = list(zip(texts, metas)) |
|
string = '' |
|
for text, meta in combined_data: |
|
# Do something with text and meta |
|
# For example, print them |
|
string += f'\n\nFrån förhör med {meta["person"]} {meta["date"]}:'.upper() |
|
string += f'\n{text}\n\n' |
|
|
|
prompt = f'''Svara på frågan: {user_input}\n |
|
Använd endast informationen nedan:\n |
|
{string}\n |
|
Skriv utförligt på svenska och var noga med detaljer som namn, plats och datum. |
|
Får gärna med information från alla fem förhör om det är relevant.\n |
|
{user_input}''' |
|
answer = st.session_state.llm.generate(prompt) |
|
return prompt |
|
|
|
st.set_page_config( |
|
page_title="Malå", |
|
) |
|
|
|
# Should not be reseted every run. |
|
if "llm" not in st.session_state: |
|
st.session_state.llm = LLM(chat=True, system_prompt='Du är assistent åt en journalist som går igenom förhör i en förundersökning. Svara bara utifrån den information du får. Svara alltid på svenska!') |
|
|
|
# Should be reset every run. |
|
llm_checker = LLM(chat=True) |
|
chroma = ChromaDB() |
|
|
|
|
|
# Initialize chat history |
|
if "messages" not in st.session_state: |
|
st.session_state.messages = [] |
|
|
|
# Display chat messages from history on app rerun |
|
for message in st.session_state.messages: |
|
with st.chat_message(message["role"]): |
|
st.markdown(message["content"]) |
|
|
|
# Accept user input |
|
if user_input := st.chat_input("Fråga något om förhören."): |
|
print_blue(user_input) |
|
if len(st.session_state.messages) > 1: |
|
history = '' |
|
for message in st.session_state.messages: |
|
history += f"{message['role']}: {message['content']}\n'" |
|
prompt = f'En användare har ställt frågan "{user_input}" och här är chatthistoriken mellan användaren och en assistent:\n{history}\n\nVerkar "{user_input}" vara en uppföljningfråga eller en fristående fråga? Svara ENDAST med "uppföljning" eller "fristående".' |
|
chat_completion = llm_checker.generate(prompt, stream=False) |
|
answer = chat_completion.choices[0].message.content |
|
|
|
print_red(answer) |
|
if 'uppföljning' in answer: |
|
prompt=f'Använd historiken till att omformulera "{user_input}" till en helt fristående fråga. Frågan ska användas för att hitta information i förhören.' |
|
chat_completion = llm_checker.generate(prompt, stream=False) |
|
question2chroma = chat_completion.choices[0].message.content |
|
if 'fristående' in answer: |
|
question2chroma = user_input |
|
|
|
if 'None' in answer: |
|
question2chroma = user_input |
|
|
|
print_yellow(question2chroma) |
|
|
|
else: |
|
question2chroma = user_input |
|
# Add user message to chat history |
|
st.session_state.messages.append({"role": "user", "content": user_input}) |
|
# Display user message in chat message container |
|
with st.chat_message("user"): |
|
st.markdown(user_input) |
|
|
|
# Display assistant response in chat message container |
|
with st.chat_message("assistant"): |
|
docs = get_docs(question2chroma) |
|
prompt = generate_prompt(user_input, docs) |
|
stream = st.session_state.llm.generate(prompt) |
|
response = st.write_stream(stream) |
|
st.session_state.llm.messages.append({'role': 'assistant', 'content': response}) |
|
st.session_state.messages.append({"role": "assistant", "content": response}) |
|
print() |
|
|
|
|
|
|