Spaces:
Build error
Build error
File size: 1,868 Bytes
6366d73 5ee695b 3f678f5 79869bb 59ee495 79869bb 5ee695b bb62c99 f5bd67a 2d6f09a bb62c99 79869bb 2d6f09a 3f678f5 47b963e 3f678f5 47b963e 3f678f5 f907830 12d32eb 9c9f7dd 12d32eb 9c9f7dd 12d32eb 3f678f5 12d32eb 9c9f7dd bb62c99 f907830 bb62c99 3f678f5 f907830 9c9f7dd 12d32eb 671ecd9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | import os
import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
def load_model():
model_id = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"
access_token = os.getenv("hf_mistral_token")
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True, token=access_token)
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype="float16"
)
model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=quant_config,
device_map="auto",
token=access_token
)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
return pipe
def main():
st.title("ChatGPT-Clone")
if "generator" not in st.session_state:
with st.spinner("Loading model..."):
st.session_state.generator = load_model()
if "messages" not in st.session_state:
st.session_state.messages = []
for msg in st.session_state.messages:
with st.chat_message(msg["role"]):
st.markdown(msg["content"])
if prompt := st.chat_input("Ask anything..."):
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
with st.spinner("Thinking..."):
result = st.session_state.generator(
prompt,
max_new_tokens=512,
temperature=0.7,
do_sample=True,
)[0]["generated_text"]
st.markdown(result)
st.session_state.messages.append({"role": "assistant", "content": result})
if __name__ == "__main__":
main()
|