Update main.py

main.py

@@ -4,45 +4,24 @@ from pydantic import BaseModel  # Import BaseModel from pydantic for data valida
 from huggingface_hub import InferenceClient  # Import InferenceClient from huggingface_hub
 import uvicorn  # Import uvicorn for running the FastAPI application
 
-app = FastAPI()  # Create a FastAPI instance
+app = FastAPI(HF_API_TOKEN)  # 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/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"
+fallbacks = ["mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mistral-7B-Instruct-v0.1"]
 
 # Define the data model for the request body
 class Item(BaseModel):
-    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
-
-
+    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
 
 # Function to generate the response JSON
 def generate_response_json(item, output, tokens, model_name):
@@ -115,7 +94,7 @@ async def generate_text(item: Item = None):
         )
 
         tokens = 0
-        client = InferenceClient(primary, token=HF_API_KEY)  # Add API key here
+        client = InferenceClient(primary)
         stream = client.text_generation(input_, **generate_kwargs, stream=True, details=True, return_full_text=True)
         output = ""
         for response in stream:
@@ -132,7 +111,7 @@ async def generate_text(item: Item = None):
         
         for model in fallbacks:
             try:
-                client = InferenceClient(model, token=HF_API_KEY)  # Add API key here for fallback models
+                client = InferenceClient(model)
                 stream = client.text_generation(input_, **generate_kwargs, stream=True, details=True, return_full_text=True)
                 output = ""
                 for response in stream:
@@ -146,10 +125,9 @@ async def generate_text(item: Item = None):
 
         raise HTTPException(status_code=500, detail=error)
 
-# Show online status
-@app.get("/")
-def root():
-    return {"status": "Sebari-chan is online!"}
+    if "KEY" in os.environ:
+        if item.key != os.environ["KEY"]:
+            raise HTTPException(status_code=401, detail="Valid key is required.")
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)