chat-with-data / model_loader.py
Alimubariz124's picture
Update model_loader.py
e685ba4 verified
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