Update app.py

app.py

@@ -24,6 +24,11 @@ try:
         max_length: int = 1024
         temperature: float = 0.7
 
+    class TokenUsage(BaseModel):
+        prompt_tokens: int
+        completion_tokens: int
+        total_tokens: int
+
     class GenerateResponse(BaseModel):
         generated_text: str
 
@@ -38,14 +43,17 @@ try:
             formatted_prompt += f"<|system|>\n{request.system_prompt}</s>\n"
             for message in request.messages:
                 if message.role == "system":
-                    formatted_prompt += f"<|system|>\n{message.content}</s>\n"
+                    formatted_prompt += f"<system>\n{message.content}\n</system>\n"
                 elif message.role == "user":
-                    formatted_prompt += f"<|user|>\n{message.content}</s>\n"
+                    formatted_prompt += f"<user>\n{message.content}\n</user>\n"
                 elif message.role == "assistant":
-                    formatted_prompt += f"<|assistant|>\n{message.content}</s>\n"
+                    formatted_prompt += f"<assistant>\n{message.content}\n</assistant>\n"
             
             # Add final assistant prefix for generation
-            formatted_prompt += "<|assistant|>\n"
+            formatted_prompt += "<assistant>\n"
+
+            # Count tokens in the prompt
+            prompt_tokens = len(tokenizer.encode(formatted_prompt))
 
             output = pipe(
                 formatted_prompt,
@@ -55,14 +63,26 @@ try:
             )[0]['generated_text']
 
             # Extract only the newly generated assistant response
-            response_text = output.split("<|assistant|>\n")[-1].split("</s>")[0]
+            response_text = output.split("<assistant>\n")[-1].split("</assistant>")[0]
+
+            # Count tokens in the completion
+            full_output_tokens = len(tokenizer.encode(output))
+            completion_tokens = full_output_tokens - prompt_tokens
+
+            usage = TokenUsage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens
+            )
         
-            return GenerateResponse(generated_text=response_text)
+            return GenerateResponse(generated_text=response_text,usage=usage)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
         #try:
             #output = pipe(request.prompt)[0]['generated_text']
             #return GenerateResponse(generated_text=output)
         #except Exception as e:
-            #raise HTTPException(status_code=500, detail=str(e))
+            #
 
 except Exception as e:
     print(f"Error: {e}")