wehub-resource-sync
55 lines
1.4 KiB
Python
55 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Third Party
|
|
from transformers import AutoTokenizer
|
|
import chat_session
|
|
import streamlit as st
|
|
|
|
# Change the following variables as needed
|
|
|
|
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
|
|
PORT = 8000
|
|
|
|
|
|
@st.cache_resource
|
|
def get_tokenizer():
|
|
global MODEL_NAME
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
|
return tokenizer
|
|
|
|
|
|
tokenizer = get_tokenizer()
|
|
|
|
|
|
@st.cache_data
|
|
def read_context() -> str:
|
|
context_file = "ffmpeg.txt"
|
|
with open(context_file, "r") as f:
|
|
context_text = f.read()
|
|
return context_text
|
|
|
|
|
|
context = read_context()
|
|
|
|
container = st.container(border=True)
|
|
|
|
with st.sidebar:
|
|
session = chat_session.ChatSession(PORT)
|
|
|
|
system_prompt = st.text_area(
|
|
"System prompt:",
|
|
"You are a helpful assistant. I will now give you a document and "
|
|
"please answer my question afterwards based on the content in document",
|
|
)
|
|
|
|
session.set_context([system_prompt] + [context])
|
|
num_tokens = tokenizer.encode(session.get_context())
|
|
container.header(
|
|
f"The context given to LLM: ({len(num_tokens)} tokens)", divider="grey"
|
|
)
|
|
container.text(session.get_context())
|
|
|
|
messages = st.container(height=400)
|
|
if prompt := st.chat_input("Type the question here"):
|
|
messages.chat_message("user").write(prompt)
|
|
messages.chat_message("assistant").write_stream(session.chat(prompt))
|