pytorch

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 
 import transformers
-from transformers import AutoTokenizer, TFAutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import tensorflow as tf
 
 SAVED_CHECKPOINT = 'mikegarts/distilgpt2-erichmariaremarque'
@@ -9,7 +9,7 @@ MIN_WORDS = 80
 
 
 def get_model():
-    model = TFAutoModelForCausalLM.from_pretrained(SAVED_CHECKPOINT)
+    model = AutoModelForCausalLM.from_pretrained(SAVED_CHECKPOINT)
     tokenizer = AutoTokenizer.from_pretrained(SAVED_CHECKPOINT)
     return model, tokenizer
 
@@ -18,14 +18,15 @@ def generate(prompt):
     model, tokenizer = get_model()
 
     input_context = prompt
-    input_ids = tokenizer.encode(input_context, return_tensors="tf")
+    input_ids = tokenizer.encode(input_context, return_tensors="pt").to('cuda')
 
     outputs = model.generate(
-        input_ids=input_ids,
-        max_length=MIN_WORDS,
-        temperature=0.7,
-        num_return_sequences=1,
-        do_sample=True
+        input_ids=input_ids, 
+        max_length=100, 
+        temperature=0.7, 
+        num_return_sequences=3, 
+        do_sample=True,
+        # forced_eos_token_id=tokenizer.encode('.')[0]
     )
 
     return tokenizer.decode(outputs[0], skip_special_tokens=True).rsplit('.', 1)[0] + '.'