Update main.py

main.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import uvicorn
+import re
 from typing import List
 
 app = FastAPI()
@@ -42,6 +43,11 @@ def generate(item: Item):
 
     for response in stream:
         output += response.token.text
+
+    # Remove unwanted sequences or patterns (e.g., <s>, [/INST], etc.)
+    output = re.sub(r'<[^>]+>', '', output)  # Remove any HTML-like tags
+    output = re.sub(r'\s+', ' ', output).strip()  # Clean up extra whitespace
+
     return output
 
 @app.get("/generate/")
@@ -63,4 +69,4 @@ async def generate_text(
         top_p=top_p,
         repetition_penalty=repetition_penalty
     )
-    return {"response": generate(item)}
+    return {"response": generate(item)}