BoxzDev commited on
Commit
7354ecc
·
verified ·
1 Parent(s): c7f61bc

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +10 -3
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)