|
import streamlit as st |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
def main(): |
|
st.title("Chatbot with Hugging Face Model") |
|
|
|
model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1" |
|
tokenizer = AutoTokenizer.from_pretrained(model_id) |
|
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto") |
|
|
|
user_input = st.text_input("User Input:", "What is your favourite condiment?") |
|
|
|
if st.button("Generate Response"): |
|
messages = [ |
|
{"role": "user", "content": user_input}, |
|
{"role": "assistant", "content": "Placeholder assistant message"} |
|
] |
|
|
|
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda") |
|
|
|
outputs = model.generate(inputs, max_new_tokens=20) |
|
response = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
st.text_area("Assistant's Response:", response) |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|