Update main.py

main.py

@@ -9,7 +9,7 @@ import replicate
 
 class Item(BaseModel):
     url: str
-    percentage: int
+    max_tokens: int
 
 app = FastAPI()
 
@@ -71,7 +71,8 @@ async def root(item: Item):
         def event_generator():
             for event in replicate.stream("snowflake/snowflake-arctic-instruct", input={
                 "prompt": "summarize this following news article:" + article,
-                "temperature": 0.2
+                "temperature": 0.2,
+                "max_new_tokens" : item.max_tokens
             }):
                 # Yield the event as a string
                 yield str(event)