Spaces:
Paused
Paused
File size: 2,731 Bytes
580cd22 861540a 580cd22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """
model.py - Inference & Model Orchestration Layer
This module encapsulates the logic for loading and executing high-parameter LLMs.
It manages the transition from raw user input to structured data by coordinating
the chat model, embedding model, and Pydantic-based output validation.
"""
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.output_parsers import PydanticOutputParser
from schema import PokemonEmeraldResponse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
CHAT_MODEL_ID = "mistralai/Mistral-Small-Instruct-2409"
EMBEDDING_MODEL_ID = "nomic-ai/nomic-embed-text-v1.5"
def load_chat_model():
"""
Bootstrap the tokenizer and chat model for the RAG pipeline.
Returns: (model, tokenizer) tuple.
"""
print(f"Loading {CHAT_MODEL_ID} model...")
tokenizer = AutoTokenizer.from_pretrained(CHAT_MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
CHAT_MODEL_ID,
torch_dtype="auto",
device_map="auto"
)
print("Model loaded!")
return model, tokenizer
def load_embedding_model():
"""
Bootstrap the HuggingFaceEmbeddings model for the RAG pipeline.
Returns: An instance of HuggingFaceEmbeddings configured with EMBEDDING_MODEL_ID.
"""
print(f"Loading {EMBEDDING_MODEL_ID} model...")
embedding_model = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_ID)
print("Model loaded!")
return embedding_model
def generate_answer(chat_model, tokenizer, prompt, max_retries=3):
"""
Executes the core RAG generation loop with automated schema validation.
This function manages the transition from raw model output to structured JSON.
It includes a retry mechanism to handle edge-case parsing failures, ensuring
the final response adheres to the PokemonEmeraldResponse Pydantic schema.
"""
print("Generating chat response...")
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
inputs_size = len(inputs.input_ids[0])
parser = PydanticOutputParser(pydantic_object=PokemonEmeraldResponse)
attempt = 0
while attempt < max_retries:
try:
response = chat_model.generate(
**inputs,
max_new_tokens=256,
num_beams=3,
early_stopping=True,
repetition_penalty=1.2,
length_penalty=1,
)
output = tokenizer.decode(response[0][inputs_size:], skip_special_tokens=True)
return parser.parse(output)
except Exception as e:
print(f"Chat model response generation failed, attempt {attempt}.")
attempt += 1
print(f"Parsing failed. Output: {output}") |