ai-forever/rugpt3small_based_on_gpt2
Browse files
main.py
CHANGED
|
@@ -14,7 +14,8 @@ class RequestData(BaseModel):
|
|
| 14 |
prompt: str
|
| 15 |
max_tokens: int = 50
|
| 16 |
|
| 17 |
-
MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
|
|
|
|
| 18 |
|
| 19 |
try:
|
| 20 |
# Загрузка модели с явным указанием device_map
|
|
|
|
| 14 |
prompt: str
|
| 15 |
max_tokens: int = 50
|
| 16 |
|
| 17 |
+
#MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
|
| 18 |
+
MODEL_NAME = "ai-forever/rugpt3small_based_on_gpt2"
|
| 19 |
|
| 20 |
try:
|
| 21 |
# Загрузка модели с явным указанием device_map
|