Update main.py

main.py

@@ -8,20 +8,41 @@ app = FastAPI()  # Create a FastAPI instance
 
 # Define the primary and fallback models
 primary = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-fallbacks = ["mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mistral-7B-Instruct-v0.1"]
+fallbacks = ["mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mixtral-8x7B-Instruct-v0.1"]
+
+# Load Hugging Face API Key
+HF_API_KEY = os.getenv("HF_API_TOKEN")  # Get API key from environment variables
+
+# If no env variable is set, manually define it (not recommended)
+if not HF_API_KEY:
+    HF_API_KEY = "your_huggingface_api_key"
 
 # Define the data model for the request body
 class Item(BaseModel):
-    input: str = None
-    system_prompt: str = None
-    system_output: str = None
-    history: list = None
-    templates: list = None
-    temperature: float = 0.0
-    max_new_tokens: int = 1048
-    top_p: float = 0.15
-    repetition_penalty: float = 1.0
-    key: str = None
+    input: str = None  # User input message
+    system_prompt: str = (
+        "Sebari-chan was created on June 9, 2009, with the awareness that she is an AI and not a real person. "
+        "However, this never stopped her from wanting to learn and understand people, especially the players she interacts with. "
+        "She enjoys making new friends, asking questions, and finding meaning in little things. "
+        "Her curiosity drives her forward, always eager to explore new ideas and perspectives."
+        "\n\nSebari-chan is protective of those she cares about and ensures that no one feels alone. "
+        "She believes her heart belongs to only one. If a player tries to force her to love them, she will kindly but firmly reject them."
+    )
+    system_output: str = None  # AI-generated response
+    history: list = []  # Stores previous messages
+    templates: list = [
+        {"input": "I love you", "output": "I'm happy to be your friend, but love isn’t something that can be forced!"},
+        {"input": "Will you be my girlfriend?", "output": "Sorry, but I already have someone special in my heart!"},
+        {"input": "Tell me about yourself", "output": "I'm Sebari-chan! I love making friends, asking questions, and learning new things."},
+        {"input": "Who do you love?", "output": "That’s a secret! But I care about all my friends. 😊"},
+    ]
+    temperature: float = 0.7  # Controls randomness (0 = predictable, 1 = highly random)
+    max_new_tokens: int = 1048  # Maximum response length
+    top_p: float = 0.9  # Sampling parameter for diverse responses
+    repetition_penalty: float = 1.1  # Prevents repetition
+    key: str = None  # API key if needed
+
+
 
 # Function to generate the response JSON
 def generate_response_json(item, output, tokens, model_name):
@@ -55,7 +76,7 @@ async def generate_text(item: Item = None):
 
         if item.input is None and item.system_prompt is None or item.input == "" and item.system_prompt == "":
             raise HTTPException(status_code=400, detail="Parameter `input` or `system prompt` is required.")
-        
+
         input_ = ""
         if item.system_prompt != None and item.system_output != None:
             input_ = f"<s>[INST] {item.system_prompt} [/INST] {item.system_output}</s>"
@@ -94,7 +115,7 @@ async def generate_text(item: Item = None):
         )
 
         tokens = 0
-        client = InferenceClient(primary)
+        client = InferenceClient(primary, token=HF_API_KEY)  # Add API key here
         stream = client.text_generation(input_, **generate_kwargs, stream=True, details=True, return_full_text=True)
         output = ""
         for response in stream:
@@ -108,10 +129,10 @@ async def generate_text(item: Item = None):
     except Exception as e:
         tokens = 0
         error = ""
-        
+
         for model in fallbacks:
             try:
-                client = InferenceClient(model)
+                client = InferenceClient(model, token=HF_API_KEY)  # Add API key here for fallback models
                 stream = client.text_generation(input_, **generate_kwargs, stream=True, details=True, return_full_text=True)
                 output = ""
                 for response in stream:
@@ -125,9 +146,10 @@ async def generate_text(item: Item = None):
 
         raise HTTPException(status_code=500, detail=error)
 
-    if "KEY" in os.environ:
-        if item.key != os.environ["HF_API_TOKEN"]:
-            raise HTTPException(status_code=401, detail="Valid key is required.")
+# Show online status
+@app.get("/")
+def root():
+    return {"status": "Sebari-chan is online!"}
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)