gemma2

main.py

@@ -10,8 +10,8 @@ import torch
 
 app = FastAPI()
 
-# Initialize the InferenceClient with your model
-client = InferenceClient("meta-llama/Llama-2-7b-chat")
+# Initialize the InferenceClient with the Gemma-7b model
+client = InferenceClient("google/gemma-7b")
 
 class Item(BaseModel):
     prompt: str
@@ -23,12 +23,16 @@ class Item(BaseModel):
     repetition_penalty: float = 1.0
 
 def format_prompt(message, history):
-    # Simple structure: alternating lines of dialogue, no special tokens unless specified by the model documentation
-    conversation = ""
-    for user_prompt, bot_response in history:
-        conversation += f"User: {user_prompt}\nBot: {bot_response}\n"
-    conversation += f"User: {message}"
-    return conversation
+    prompt = "<bos>"
+    # Add history to the prompt if there's any
+    if history:
+        for entry in history:
+            role = "user" if entry['role'] == "user" else "model"
+            prompt += f"<start_of_turn>{role}\n{entry['content']}<end_of_turn>"
+    # Add the current message
+    prompt += f"<start_of_turn>user\n{message}<end_of_turn><start_of_turn>model\n"
+    return prompt
+