yt-video-to-summary / models /flan_t5_xl.py
zman1x1's picture
add models
bfadc34 unverified
raw
history blame
490 Bytes
import os, torch, accelerate
from langchain.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, AutoModelForCausalLM
model_id = 'google/flan-t5-large'
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id, load_in_8bit=True)
pipe = pipeline(
"text2text-generation",
model=model,
tokenizer=tokenizer,
max_length=512,
)
local_llm = HuggingFacePipeline(Pipeline=pipe)