modified enpoints

app.py

@@ -1,5 +1,15 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
+from app.main import app as main_app
+
+
 app = FastAPI()
+
+
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
+
+# Mount the main app
+app.mount("/", main_app)
+
+

app/chains.py

@@ -6,7 +6,7 @@ from prompts import (
     raw_prompt,
     raw_prompt_formatted,
     format_context,
-    # tokenizer
+    tokenizer
 )
 from data_indexing import DataIndexer
 
@@ -16,7 +16,7 @@ llm = HuggingFaceEndpoint(
     model="meta-llama/Llama-3.1-8B-Instruct",
     huggingfacehub_api_token=os.environ['HF_TOKEN'],
     max_new_tokens=512,
-    # stop_sequences=[tokenizer.eos_token],
+    stop_sequences=[tokenizer.eos_token],
     streaming=True,
 )
 

app/prompts.py

@@ -1,12 +1,29 @@
 from langchain_core.prompts import PromptTemplate
 from typing import List
 import models
+from transformers import AutoTokenizer
+
+model_name = "meta-llama/Llama-3.1-8B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
 
 
 def format_prompt(prompt) -> PromptTemplate:
     # TODO: format the input prompt by using the model specific instruction template
     # TODO: return a langchain PromptTemplate
-    return PromptTemplate.from_template(prompt)
+    # Test with a simple question
+    chat = [
+        {"role": "system", "content": "You are a helpful AI assistant."},
+        {"role": "user", "content": prompt},
+    ]
+    
+    formatted_prompt = tokenizer.apply_chat_template(
+        chat, 
+        tokenize=False, 
+        add_generation_prompt=True
+    )
+
+    return PromptTemplate.from_template(formatted_prompt)
 
 def format_chat_history(messages: List[models.Message]):
     # TODO:  implement format_chat_history to format 
@@ -35,7 +52,7 @@ standalone_prompt: str = None
 rag_prompt: str = None
 
 # TODO: create raw_prompt_formatted by using format_prompt
-raw_prompt_formatted = None
+raw_prompt_formatted = format_prompt(raw_prompt)
 raw_prompt = PromptTemplate.from_template(raw_prompt)
 
 # TODO: use format_prompt to create history_prompt_formatted
@@ -43,9 +60,4 @@ history_prompt_formatted: PromptTemplate = None
 # TODO: use format_prompt to create standalone_prompt_formatted
 standalone_prompt_formatted: PromptTemplate = None
 # TODO: use format_prompt to create rag_prompt_formatted
-rag_prompt_formatted: PromptTemplate = None
-
-
-
-
-
+rag_prompt_formatted: PromptTemplate = None