AutoModel

main.py

@@ -4,15 +4,18 @@ os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf/transformers"
 
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, AutoModel
 from fastapi.responses import StreamingResponse
 import threading
 
 app = FastAPI()
 
-model_id = "GEB-AGI/geb-1.3b"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
+# model_id = "GEB-AGI/geb-1.3b"
+# tokenizer = AutoTokenizer.from_pretrained(model_id)
+# model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
+model = AutoModel.from_pretrained("GEB-AGI/geb-1.3b", trust_remote_code=True).bfloat16().cuda()
+tokenizer = AutoTokenizer.from_pretrained("GEB-AGI/geb-1.3b", trust_remote_code=True)
+
 
 class ChatRequest(BaseModel):
     message: str