File size: 427 Bytes
a361ca0
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
from .retriever import CustomRetriever
from .generator import CustomGenerator


class MyRagModel:
    def __init__(self, retriever: CustomRetriever, generator: CustomGenerator):
        self.retriever = retriever
        self.generator = generator

    def get_answer(self, query):
        context = self.retriever.retrieve(query)
        response = self.generator.generate(query, context)
        return response