Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 from typing import Dict, List, Optional
+from dataclasses import dataclass, field
 from transformers import (
     AutoConfig,
     AutoModelForSequenceClassification,
@@ -10,6 +11,17 @@ from transformers import (
     TrainingArguments,
 )
 
+@dataclass
+class ModelArguments:
+    model_name: str = field(
+        default="gpt2",
+        metadata={"help": "The name of the pretrained model to use for text generation."}
+    )
+    max_tokens: int = field(
+        default=50,
+        metadata={"help": "The maximum number of tokens to generate in the response."}
+    )
+
 class MockOpenAI:
     """
     A mock implementation of OpenAI's API using Hugging Face's pipeline for text generation.
@@ -96,11 +108,9 @@ class MockOpenAI:
 
 # Example usage
 if __name__ == "__main__":
-    parser = HfArgumentParser(description="Mock OpenAI API using Hugging Face's pipeline for text generation.")
-    parser.add_argument("--model_name", default="gpt2", help="The name of the pretrained model to use for text generation.")
-    parser.add_argument("--max_tokens", type=int, default=50, help="The maximum number of tokens to generate in the response.")
-    args = parser.parse_args()
-    client = MockOpenAI(model_name=args.model_name, max_tokens=args.max_tokens)
+    parser = HfArgumentParser((ModelArguments,))
+    model_args = parser.parse_args_into_dataclasses()[0]
+    client = MockOpenAI(model_name=model_args.model_name, max_tokens=model_args.max_tokens)
     chat_completion = client.chat.Completions().create(
         messages=[
             {
@@ -114,4 +124,4 @@ if __name__ == "__main__":
         ]
     )
 
-    print(chat_completion)
+    print(chat_completion)