Update app.py

app.py

@@ -235,7 +235,7 @@ import time
 from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
-from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings
+from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, PromptTemplate, Settings
 from llama_index.llms.huggingface import HuggingFaceLLM
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from pydantic import BaseModel
@@ -244,7 +244,7 @@ import datetime
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.templating import Jinja2Templates
 from simple_salesforce import Salesforce, SalesforceLogin
-from transformers import AutoModelForSeq2SeqLM
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 # Define Pydantic model for incoming request body
 class MessageRequest(BaseModel):
@@ -288,6 +288,7 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="static")
 
 # Configure Llama index settings
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
 Settings.llm = HuggingFaceLLM(
     model_name="google/flan-t5-small",
     tokenizer_name="google/flan-t5-small",
@@ -295,6 +296,7 @@ Settings.llm = HuggingFaceLLM(
     max_new_tokens=256,
     generate_kwargs={"temperature": 0.1, "do_sample": True},
     model=AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small"),
+    tokenizer=tokenizer,
     device_map="auto"  # Automatically use GPU if available, else CPU
 )
 Settings.embed_model = HuggingFaceEmbedding(
@@ -341,18 +343,15 @@ def split_name(full_name):
 initialize()  # Run initialization tasks
 
 def handle_query(query):
-    chat_text_qa_msgs = [
-        (
-            "user",
-            """
-            You are Clara, a Redfernstech chatbot. Provide accurate, concise answers (10-15 words) based on company data.
-            Context: {context_str}
-            Question: {query_str}
-            Answer:
-            """
-        )
-    ]
-    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
+    # Custom prompt template for flan-t5-small (no chat template)
+    text_qa_template = PromptTemplate(
+        """
+        You are Clara, a Redfernstech chatbot. Provide accurate, concise answers (10-15 words) based on company data.
+        Context: {context_str}
+        Question: {query_str}
+        Answer:
+        """
+    )
     
     storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
     index = load_index_from_storage(storage_context)
@@ -361,7 +360,7 @@ def handle_query(query):
         if past_query.strip():
             context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"
     
-    query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str)
+    query_engine = index.as_query_engine(text_qa_template=text_qa_template)
     answer = query_engine.query(query)
 
     if hasattr(answer, "response"):