Update main.py

main.py

@@ -10,6 +10,13 @@ app = FastAPI()  # Create a FastAPI instance
 primary = "HuggingFaceH4/zephyr-7b-beta"
 fallbacks = ["mistralai/Mistral-7B-Instruct-v0.2", "HuggingFaceH4/zephyr-7b-beta"]
 
+# Load Hugging Face API Key
+HF_API_KEY = os.getenv("")  # 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
@@ -94,7 +101,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:
@@ -111,7 +118,7 @@ async def generate_text(item: Item = None):
         
         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:
@@ -130,4 +137,4 @@ async def generate_text(item: Item = None):
             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)