spuun commited on
Commit
be6fdd2
·
1 Parent(s): 3bd4075

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -0
app.py ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoTokenizer
3
+ from petals import AutoDistributedModelForCausalLM
4
+
5
+ model_name = "petals-team/StableBeluga2"
6
+
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, add_bos_token=False)
8
+ model = AutoDistributedModelForCausalLM.from_pretrained(model_name)
9
+
10
+ import gradio as gr
11
+
12
+ def generate(input, params):
13
+ tokenized = tokenizer(input, return_tensors="pt")["input_ids"]
14
+ outputs = model.generate(inputs, max_new_tokens=80, do_sample=True, temperature=0.9)
15
+ return tokenizer.decode(outputs[0]).replace("</s>", "");
16
+
17
+ iface = gr.Interface(fn=generate, inputs="text", outputs="text")
18
+ iface.launch()