File size: 496 Bytes
37f30a1
e685ba4
37f30a1
 
 
 
 
2e1a242
37f30a1
2e1a242
 
37f30a1
1
2
3
4
5
6
7
8
9
10
11
12
13
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

def load_embedding_model():
    return SentenceTransformer("all-MiniLM-L6-v2")

def load_llm():
    model_name = "google/flan-t5-base"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer, max_new_tokens=200)
    return pipe