use larger model

app.py

@@ -1,5 +1,4 @@
 from indexer import (
-    load_raw_dataset,
     create_vector_database,
     get_llm,
     get_prompt_template,
@@ -24,14 +23,15 @@ class CustomRAG:
 
     def run(self, query):
         retriever = self.vector_db.as_retriever(search_kwargs={"k": 3})
-        contexts = retriever.get_relevant_documents(query)
+        contexts = retriever.invoke(query)
         formatted_context = format_contexts(contexts)
         prompt = self.prompt_template.format(context=formatted_context, question=query)
         return self.llm.invoke(prompt), contexts
 
 
 def answer_question(query):
-    llm = get_llm("google/flan-t5-base")
+    # llm = get_llm("google/flan-t5-base")
+    llm = get_llm("FreedomIntelligence/HuatuoGPT-o1-7B")
     vector_database = create_vector_database("sentence-transformers/all-MiniLM-L6-v2")
     prompt_template = get_prompt_template()
     rag = CustomRAG(

indexer.py

@@ -7,6 +7,7 @@ from transformers import (
     AutoTokenizer,
     pipeline,
     AutoModelForSeq2SeqLM,
+    AutoModelForCausalLM,
 )
 from langchain.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
@@ -16,6 +17,7 @@ api_key = os.environ["PINECONE_API_KEY"]
 
 from langchain_pinecone import PineconeVectorStore
 
+
 def load_raw_dataset():
     dataset = load_dataset("lavita/ChatDoctor-HealthCareMagic-100k")
 
@@ -47,16 +49,28 @@ def create_vector_database(model_name):
 
 def get_llm(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(
-        "google/flan-t5-base", torch_dtype="auto", device_map="auto"
+    # model = AutoModelForSeq2SeqLM.from_pretrained(
+    #     "google/flan-t5-base", torch_dtype="auto", device_map="auto"
+    # )
+
+    # pipe = pipeline(
+    #     "text2text-generation",
+    #     model=model,
+    #     tokenizer=tokenizer,
+    #     max_new_tokens=512,
+    #     temperature=1,
+    #     do_sample=True,
+    # )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, torch_dtype="auto", device_map="auto"
     )
-
     pipe = pipeline(
-        "text2text-generation",
+        "text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_new_tokens=512,
-        temperature=1,
+        max_new_tokens=1024,
+        temperature=0.7,
         do_sample=True,
     )