Pippoz commited on
Commit
7d5081f
1 Parent(s): da92359

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
+ import time
3
+ from datetime import datetime
4
+ from transformers import AutoModelForCausalLM, AutoTokenizer
5
+ import torch
6
+
7
+ @st.cache(allow_output_mutation=True)
8
+ def opt_model(prompt, num_sequences = 1, max_length = 50):
9
+ model = AutoModelForCausalLM.from_pretrained("facebook/opt-30b", torch_dtype=torch.float16).cuda()
10
+ tokenizer = AutoTokenizer.from_pretrained("facebook/opt-30b", use_fast=False)
11
+ input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
12
+ generated_ids = model.generate(input_ids, num_return_sequences=num_sequences, max_length=max_length)
13
+ answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
14
+ return answer
15
+
16
+ prompt= st.text_area('Your prompt here',
17
+ '''Hello, I'm am conscious and''')
18
+
19
+ answer = opt_model(prompt)
20
+ #lst = ['ciao come stai sjfsbd dfhsdf fuahfuf feuhfu wefwu ']
21
+ lst = ' '.join(answer)
22
+
23
+ t = st.empty()
24
+ for i in range(len(lst)):
25
+ t.markdown("### %s..." % lst[0:i])
26
+ time.sleep(0.04)