Tech-Meld commited on
Commit
020026a
·
verified ·
1 Parent(s): 7a95668

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +26 -0
app.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+
4
+ @st.cache(allow_output_mutation=True)
5
+ def load_model():
6
+ model_id = "Tech-Meld/Hajax_Chat_1.0-Q3_K_S-GGUF"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
8
+ model = AutoModelForCausalLM.from_pretrained(model_id)
9
+ return model, tokenizer
10
+
11
+ def get_response(input_text, model, tokenizer):
12
+ inputs = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')
13
+ outputs = model.generate(inputs, max_length=1000, pad_token_id=tokenizer.eos_token_id)
14
+ response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
15
+ return response
16
+
17
+ def main():
18
+ model, tokenizer = load_model()
19
+ st.title("Chat with AI")
20
+ input_text = st.text_input("You: ", "")
21
+ if st.button("Send"):
22
+ response = get_response(input_text, model, tokenizer)
23
+ st.text_area("AI: ", response)
24
+
25
+ if __name__ == "__main__":
26
+ main()