update model

app.py

@@ -1,17 +1,12 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 import torch
 
-model_name = "EleutherAI/pythia-1b"
-model = AutoModelForCausalLM.from_pretrained(
-  model_name,
-  cache_dir="./pythia-1b",
-)
 
-tokenizer = AutoTokenizer.from_pretrained(
-  model_name,
-  cache_dir="./pythia-1b",
-)
+
+model_name = "google/flan-t5-base"
+model = T5ForConditionalGeneration.from_pretrained(model_name)
+tokenizer = T5Tokenizer.from_pretrained(model_name)
 
 
 def predict(message, history):
@@ -23,16 +18,10 @@ def predict(message, history):
     input_ids = tokenizer.encode(message, return_tensors="pt")
     chat_history_ids = model.generate(
         input_ids,
-        min_length=20,
-        max_new_tokens=128,
+        max_length=512,
         do_sample=True,
-        top_p=0.95,
-        top_k=50,
-        temperature=0.75,
-        num_return_sequences=5,
-        pad_token_id=tokenizer.eos_token_id,
     )
-    response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+    response = tokenizer.decode(chat_history_ids[0], skip_special_tokens=True)
     yield response
 
 

requirements.txt

@@ -2,3 +2,4 @@ gradio
 transformers
 # -i https://download.pytorch.org/whl/cpu
 torch
+sentencepiece