mn9206986 commited on
Commit
0253026
·
verified ·
1 Parent(s): d0aa793

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -0
app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import torch
4
+
5
+ # Load model
6
+ model_name = "OpenAssistant/oasst-sft-3-gpt-j-3B"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForCausalLM.from_pretrained(model_name)
9
+
10
+ # Device
11
+ device = torch.device("cpu")
12
+ model.to(device)
13
+
14
+ def generate(prompt):
15
+ inputs = tokenizer(prompt, return_tensors="pt").to(device)
16
+ outputs = model.generate(
17
+ **inputs,
18
+ max_new_tokens=200,
19
+ do_sample=True,
20
+ temperature=0.7
21
+ )
22
+ text = tokenizer.decode(outputs[0], skip_special_tokens=True)
23
+ return text
24
+
25
+ # Gradio interface làm API
26
+ iface = gr.Interface(
27
+ fn=generate,
28
+ inputs=gr.Textbox(lines=5, placeholder="Nhập prompt…"),
29
+ outputs=gr.Textbox(),
30
+ title="OASST-J-3B API"
31
+ )
32
+
33
+ iface.launch(server_name="0.0.0.0", server_port=7860)