Update main.py

main.py

@@ -3,6 +3,11 @@ from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import uvicorn
 import os
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+)
 
 app = FastAPI()