Erik Sarriegui commited on
Commit
ceea4f8
·
1 Parent(s): 0ada016
Files changed (1) hide show
  1. app.py +45 -0
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
3
+ from threading import Thread
4
+
5
+ model_id = "AuriLab/gpt-bi"
6
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
7
+ model = AutoModelForCausalLM.from_pretrained(model_id)
8
+
9
+ def generar_texto(prompt):
10
+ streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
11
+
12
+ inputs = tokenizer(prompt, return_tensors="pt")
13
+
14
+ thread = Thread(target=model.generate, kwargs={
15
+ "input_ids": inputs.input_ids,
16
+ "max_new_tokens": 50,
17
+ "streamer": streamer,
18
+ "do_sample": True
19
+ })
20
+ thread.start()
21
+
22
+ output = prompt
23
+ yield output
24
+
25
+ for new_token in streamer:
26
+ output += new_token
27
+ yield output
28
+
29
+ textbox = gr.Textbox(
30
+ lines=5,
31
+ placeholder="Escribe tu prompt aquí...",
32
+ label="Generación en tiempo real"
33
+ )
34
+
35
+
36
+ interfaz = gr.Interface(
37
+ fn=generar_texto,
38
+ inputs=textbox,
39
+ outputs=textbox,
40
+ title="Generador de Texto con Streamer",
41
+ description="Escribe y los tokens aparecerán en tiempo real en la misma caja"
42
+ )
43
+
44
+ if __name__ == "__main__":
45
+ interfaz.launch()