revert to flan-5 model

app.py

@@ -30,8 +30,8 @@ class CustomRAG:
 
 
 def answer_question(query):
-    # llm = get_llm("google/flan-t5-base")
-    llm = get_llm("FreedomIntelligence/HuatuoGPT-o1-7B")
+    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

@@ -49,31 +49,31 @@ 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"
-    # )
-
-    # 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"
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        "google/flan-t5-base", torch_dtype="auto", device_map="auto"
     )
+
     pipe = pipeline(
-        "text-generation",
+        "text2text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_new_tokens=1024,
-        temperature=0.7,
+        max_new_tokens=512,
+        temperature=1,
         do_sample=True,
     )
 
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     model_name, torch_dtype="auto", device_map="auto"
+    # )
+    # pipe = pipeline(
+    #     "text-generation",
+    #     model=model,
+    #     tokenizer=tokenizer,
+    #     max_new_tokens=1024,
+    #     temperature=0.7,
+    #     do_sample=True,
+    # )
+
     llm = HuggingFacePipeline(pipeline=pipe)
     return llm