update

app.py

@@ -2,8 +2,8 @@ import gradio as gr
 from transformers import AutoTokenizer, pipeline
 import torch
 
-tokenizer1 = AutoTokenizer.from_pretrained("notexist/ttt2")
-tdk1 = pipeline('text-generation', model='notexist/ttt2', tokenizer=tokenizer1)
+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)
 
@@ -23,11 +23,14 @@ def predict(name, sl, topk, topp):
             max_length=64, 
             top_k=topk, 
             top_p=topp, 
-            num_return_sequences=1,
+            num_return_sequences=2,
             repetition_penalty=sl
         )[0]["generated_text"]
-
-        return x1[len(f"<|endoftext|>"):]+"\n\n"+x2[len(f"<|endoftext|>{new_name}\n\n"):].replace("[TEXT]", " "+new_name+" ")
+        
+        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, 
@@ -42,11 +45,14 @@ def predict(name, sl, topk, topp):
             max_length=64, 
             top_k=topk, 
             top_p=topp, 
-            num_return_sequences=1,
+            num_return_sequences=2,
             repetition_penalty=sl
         )[0]["generated_text"]
 
-        return x1[len(f"<|endoftext|>{name}\n\n"):]+"\n\n"+x2[len(f"<|endoftext|>{name}\n\n"):].replace("[TEXT]", " "+name+" ")
+        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+" ")