Spaces:
Sleeping
Sleeping
Benjamin Gonzalez
commited on
Commit
·
c778ae5
1
Parent(s):
f8fff4a
create space
Browse files
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)
|