dquisi commited on
Commit
4e5beb5
1 Parent(s): b1a4130

Create new file

Browse files
Files changed (1) hide show
  1. app.py +51 -0
app.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import *
2
+ import gradio as gr
3
+ import torch
4
+ from googletrans import Translator
5
+ model_category = AutoModelForCausalLM.from_pretrained("dquisi/story_spanish_gpt2_by_category")
6
+ tokenizer_category = AutoTokenizer.from_pretrained("dquisi/story_spanish_gpt2_by_category")
7
+ story_gen = pipeline("text-generation", "pranavpsv/gpt2-genre-story-generator")
8
+ translator = Translator()
9
+ tokenizer = AutoTokenizer.from_pretrained("dquisi/story_spanish_gpt2_v2")
10
+ model = AutoModelForCausalLM.from_pretrained("dquisi/story_spanish_gpt2_v2")
11
+ def generate_story_translate(texto,longitud=250,categoria='superheroe'):
12
+ input_translate = translator.translate(texto, src='es', dest='en')
13
+ input_translate_category = translator.translate(categoria, src='es', dest='en')
14
+ query = "<BOS> <{0}> {1}".format(input_translate_category.text, input_translate.text)
15
+ generated_text = story_gen(query, max_length=longitud,do_sample=True,repetition_penalty=1.1, temperature=1.2, top_p=0.95, top_k=50)
16
+ generated_text = generated_text[0]['generated_text']
17
+ generated_text = generated_text.split('> ')[2]
18
+ generate_translate = translator.translate(generated_text, src='en', dest='es')
19
+ return generate_translate.text
20
+ def generate_story_custom(texto,longitud=250):
21
+ query = "<BOS> <{0}>".format(texto)
22
+ input_ids = tokenizer(query, return_tensors="pt")["input_ids"]
23
+ output = model.generate(input_ids, max_length=longitud,do_sample=True,repetition_penalty=1.1, temperature=1.2, top_p=0.95, top_k=50)
24
+ return tokenizer.decode(output[0])
25
+ def generate_story_custom_category(texto,longitud=250,categoria='superheroe'):
26
+ query = "<BOS> <{0}> {1}".format(categoria, texto)
27
+ input_ids = tokenizer_category(query, return_tensors="pt")["input_ids"]
28
+ output = model_category.generate(input_ids, max_length=longitud,do_sample=True,repetition_penalty=1.1, temperature=1.2, top_p=0.95, top_k=50)
29
+ return tokenizer_category.decode(output[0])
30
+
31
+ contexto = gr.inputs.Textbox(lines=10, placeholder="Ingresar palabras claves para generar un cuento")
32
+ categoria = gr.inputs.Textbox(lines=1, placeholder="Ingresar Categoria")
33
+ longitud = gr.inputs.Slider(50, 500)
34
+ opciones = gr.inputs.CheckboxGroup(["Generar", "Generar por Categoria", "Generar 2"])
35
+ resultado = gr.outputs.HTML(label="Resultado")
36
+ def generate_storie(contexto,opciones,categoria,longitud):
37
+ resultado=""
38
+ cuentos=[]
39
+ if "Generar" in opciones:
40
+ cuentos.append(generate_story_custom(contexto,longitud))
41
+ if "Generar por Categoria" in opciones:
42
+ cuentos.append(generate_story_custom_category(contexto,longitud,categoria))
43
+ if "Generar 2" in opciones:
44
+ cuentos.append(generate_story_translate(contexto,longitud,categoria))
45
+ resultado += "<p><b>Generados:</b> "+'<b>Cuento: <b/>'.join(cuentos)+"</p>"
46
+ return resultado
47
+ iface = gr.Interface(
48
+ fn=generate_storie,
49
+ inputs=[contexto,categoria,longitud,opciones],
50
+ outputs=resultado)
51
+ iface.launch(debug=True)