NightPrince commited on
Commit
1b238db
·
verified ·
1 Parent(s): c14c48b

Create streamlit_app.py

Browse files
Files changed (1) hide show
  1. streamlit_app.py +59 -0
streamlit_app.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
3
+ import torch
4
+
5
+ st.set_page_config(page_title="OpenChat 3.5 - HF Space", page_icon="🤖")
6
+ st.title("🤖 YahyaGPT")
7
+
8
+ @st.cache_resource
9
+ def load_model():
10
+ tokenizer = AutoTokenizer.from_pretrained("openchat/openchat-3.5-1210")
11
+ model = AutoModelForCausalLM.from_pretrained(
12
+ "openchat/openchat-3.5-1210",
13
+ torch_dtype=torch.float16,
14
+ device_map="auto"
15
+ )
16
+ return tokenizer, model
17
+
18
+ tokenizer, model = load_model()
19
+
20
+ if "messages" not in st.session_state:
21
+ st.session_state.messages = []
22
+
23
+ for msg in st.session_state.messages:
24
+ st.chat_message(msg["role"]).write(msg["content"])
25
+
26
+
27
+ prompt = st.chat_input("Ask Yahya anything...")
28
+
29
+ if prompt:
30
+ st.chat_message("user").write(prompt)
31
+ st.session_state.messages.append({"role": "user", "content": prompt})
32
+
33
+ full_prompt = "You are a helpful assistant.\n"
34
+ for msg in st.session_state.messages:
35
+ role = msg["role"]
36
+ content = msg["content"]
37
+ if role == "user":
38
+ full_prompt += f"User: {content}\n"
39
+ else:
40
+ full_prompt += f"Assistant: {content}\n"
41
+ full_prompt += "Assistant:"
42
+
43
+ inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
44
+
45
+ output = model.generate(
46
+ **inputs,
47
+ max_new_tokens=300,
48
+ do_sample=True,
49
+ temperature=0.7,
50
+ top_p=0.9,
51
+ repetition_penalty=1.1,
52
+ pad_token_id=tokenizer.eos_token_id
53
+ )
54
+
55
+ response = tokenizer.decode(output[0], skip_special_tokens=True)
56
+ reply = response.split("Assistant:")[-1].strip()
57
+
58
+ st.chat_message("assistant").write(reply)
59
+ st.session_state.messages.append({"role": "assistant", "content": reply})