Spaces:
Sleeping
Sleeping
Update main.py
Browse files
main.py
CHANGED
@@ -10,6 +10,13 @@ app = FastAPI() # Create a FastAPI instance
|
|
10 |
primary = "HuggingFaceH4/zephyr-7b-beta"
|
11 |
fallbacks = ["mistralai/Mistral-7B-Instruct-v0.2", "HuggingFaceH4/zephyr-7b-beta"]
|
12 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
# Define the data model for the request body
|
14 |
class Item(BaseModel):
|
15 |
input: str = None
|
@@ -94,7 +101,7 @@ async def generate_text(item: Item = None):
|
|
94 |
)
|
95 |
|
96 |
tokens = 0
|
97 |
-
client = InferenceClient(primary)
|
98 |
stream = client.text_generation(input_, **generate_kwargs, stream=True, details=True, return_full_text=True)
|
99 |
output = ""
|
100 |
for response in stream:
|
@@ -111,7 +118,7 @@ async def generate_text(item: Item = None):
|
|
111 |
|
112 |
for model in fallbacks:
|
113 |
try:
|
114 |
-
client = InferenceClient(model)
|
115 |
stream = client.text_generation(input_, **generate_kwargs, stream=True, details=True, return_full_text=True)
|
116 |
output = ""
|
117 |
for response in stream:
|
@@ -130,4 +137,4 @@ async def generate_text(item: Item = None):
|
|
130 |
raise HTTPException(status_code=401, detail="Valid key is required.")
|
131 |
|
132 |
if __name__ == "__main__":
|
133 |
-
uvicorn.run(app, host="0.0.0.0", port=8000)
|
|
|
10 |
primary = "HuggingFaceH4/zephyr-7b-beta"
|
11 |
fallbacks = ["mistralai/Mistral-7B-Instruct-v0.2", "HuggingFaceH4/zephyr-7b-beta"]
|
12 |
|
13 |
+
# Load Hugging Face API Key
|
14 |
+
HF_API_KEY = os.getenv("") # Get API key from environment variables
|
15 |
+
|
16 |
+
# If no env variable is set, manually define it (not recommended)
|
17 |
+
if not HF_API_KEY:
|
18 |
+
HF_API_KEY = "your_huggingface_api_key"
|
19 |
+
|
20 |
# Define the data model for the request body
|
21 |
class Item(BaseModel):
|
22 |
input: str = None
|
|
|
101 |
)
|
102 |
|
103 |
tokens = 0
|
104 |
+
client = InferenceClient(primary, token=HF_API_KEY) # Add API key here
|
105 |
stream = client.text_generation(input_, **generate_kwargs, stream=True, details=True, return_full_text=True)
|
106 |
output = ""
|
107 |
for response in stream:
|
|
|
118 |
|
119 |
for model in fallbacks:
|
120 |
try:
|
121 |
+
client = InferenceClient(model, token=HF_API_KEY) # Add API key here for fallback models
|
122 |
stream = client.text_generation(input_, **generate_kwargs, stream=True, details=True, return_full_text=True)
|
123 |
output = ""
|
124 |
for response in stream:
|
|
|
137 |
raise HTTPException(status_code=401, detail="Valid key is required.")
|
138 |
|
139 |
if __name__ == "__main__":
|
140 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|