File size: 2,370 Bytes
502f071
 
 
 
d214fd2
 
b7eb50d
 
502f071
10a2445
 
7c6863d
 
 
 
 
 
 
 
b7eb50d
 
10a2445
 
 
 
d214fd2
10a2445
 
d214fd2
 
 
 
 
10a2445
7c6863d
 
 
 
 
 
 
 
 
10a2445
 
 
 
d214fd2
10a2445
 
502f071
d214fd2
 
 
 
502f071
 
10a2445
 
 
 
 
 
502f071
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import gradio as gr
from transformers import AutoTokenizer, pipeline
import torch

tokenizer1 = AutoTokenizer.from_pretrained("notexist/tttf")
tdk1 = pipeline('text-generation', model='notexist/tttf', tokenizer=tokenizer1)
tokenizer2 = AutoTokenizer.from_pretrained("notexist/ttte")
tdk2 = pipeline('text-generation', model='notexist/ttte', tokenizer=tokenizer2)

def predict(name, sl, topk, topp):
    if name == "":
        x1 = tdk1(f"<|endoftext|>",
            do_sample=True, 
            max_length=64, 
            top_k=topk, 
            top_p=topp, 
            num_return_sequences=1,
            repetition_penalty=sl
        )[0]["generated_text"]
        new_name = x1[len(f"<|endoftext|>"):x1.index("\n\n")]
        x2 = tdk2(f"<|endoftext|>{new_name}\n\n",
            do_sample=True, 
            max_length=64, 
            top_k=topk, 
            top_p=topp, 
            num_return_sequences=2,
            repetition_penalty=sl
        )[0]["generated_text"]
        
        if "[TEXT]" not in x2:
            return x1[len(f"<|endoftext|>"):]   
        else:
            return x1[len(f"<|endoftext|>"):]+"\n\n"+x2[len(f"<|endoftext|>{new_name}\n\n"):].replace("[TEXT]", " "+new_name+" ")
    else:
        x1 = tdk1(f"<|endoftext|>{name}\n\n",
            do_sample=True, 
            max_length=64, 
            top_k=topk, 
            top_p=topp, 
            num_return_sequences=1,
            repetition_penalty=sl
        )[0]["generated_text"]
        x2 = tdk2(f"<|endoftext|>{name}\n\n",
            do_sample=True, 
            max_length=64, 
            top_k=topk, 
            top_p=topp, 
            num_return_sequences=2,
            repetition_penalty=sl
        )[0]["generated_text"]

        if "[TEXT]" not in x2:
            return x1[len(f"<|endoftext|>{name}\n\n"):]
        else:
            return x1[len(f"<|endoftext|>{name}\n\n"):]+"\n\n"+x2[len(f"<|endoftext|>{name}\n\n"):].replace("[TEXT]", " "+name+" ")



iface = gr.Interface(fn=predict, inputs=["text",\
                                         gr.inputs.Slider(0, 3, default=1.1, label="repetition_penalty"),\
                                         gr.inputs.Slider(0, 100, default=75, label="top_k"),\
                                         gr.inputs.Slider(0, 1, default=0.95, label="top_p")]
                     , outputs="text")
iface.launch()