Benjamin Gonzalez commited on
Commit
c778ae5
·
1 Parent(s): f8fff4a

create space

Browse files
Files changed (1) hide show
  1. app.py +12 -0
app.py ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer, AutoModelForCausalLM
2
+ import streamlit as st
3
+
4
+ tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
5
+ model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", flash_attn=True, flash_rotary=True, fused_dense=True, device_map="cuda", trust_remote_code=True)
6
+
7
+ prompt = st.text_input("Input prompt", value="Write a detailed analogy between mathematics and a lighthouse.")
8
+ length = st.number_input("Max token length", value=200)
9
+ inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
10
+ outputs = model.generate(**inputs, max_length=length)
11
+ text = tokenizer.batch_decode(outputs)[0]
12
+ st.write(text)