Spaces:

zman1x1
/

yt-video-to-summary

Running

zman1x1 commited on Jul 9, 2023

Commit

76842ff

•

1 Parent(s): 1cce785

cpu mode

Files changed (1) hide show

models/t5_small_medium_title_generation.py CHANGED Viewed

@@ -4,7 +4,7 @@ import torch
 def t5model(prompt: str) -> str:
     tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-small-medium-title-generation")
-    model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-small-medium-title-generation", torch_dtype=torch.float16)
     inputs = tokenizer(
         ["summarize:" + prompt],
         return_tensors="pt",

 def t5model(prompt: str) -> str:
     tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-small-medium-title-generation")
+    model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-small-medium-title-generation", torch_dtype=torch.float32)
     inputs = tokenizer(
         ["summarize:" + prompt],
         return_tensors="pt",