neo

app.py

@@ -1,13 +1,23 @@
 import gradio as gr
 from transformers import pipeline, set_seed
+from transformers import GPTNeoForCausalLM, GPT2Tokenizer
 
+model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
+tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
 
-generator = pipeline('text-generation', model='EleutherAI/GPTNeoForCausalLM', return_full_text=False)
 set_seed(42)
 
 
 def Bemenet(bemenet):
-    return generator(bemenet, max_length=100, num_return_sequences=2)[1]['generated_text']
+    input_ids = tokenizer(bemenet, return_tensors="pt").input_ids
+
+    gen_tokens = model.generate(
+        input_ids,
+        do_sample=True,
+        temperature=0.9,
+        max_length=100,
+    )
+    return tokenizer.batch_decode(gen_tokens)[0]
 
 
 interface = gr.Interface(fn=Bemenet,