solar

question_paper.py

@@ -10,24 +10,27 @@ app = FastAPI()
 MODEL = None
 TOKENIZER = None
 
-
 @app.get("/")
 def llama():
     text = "Hi, my name is "
-    inputs = TOKENIZER(text, return_tensors="pt")
-    outputs = MODEL.generate(**inputs, max_new_tokens=64)
+    inputs = TOKENIZER(text, return_tensors="pt").input_ids
+    outputs = MODEL.generate(
+        inputs,
+        max_length=256,
+        pad_token_id=TOKENIZER.pad_token_id,
+        eos_token_id=TOKENIZER.eos_token_id,
+    )
     tresponse = TOKENIZER.decode(outputs[0], skip_special_tokens=True)
     print(tresponse)
 
     return tresponse
 
-
 @app.on_event("startup")
 def init_model():
     global MODEL
     global TOKENIZER
     if not MODEL:
         print("loading model")
-        TOKENIZER = AutoTokenizer.from_pretrained("Upstage/SOLAR-10.7B-v1.0")
-        MODEL = AutoModelForCausalLM.from_pretrained("Upstage/SOLAR-10.7B-v1.0", device_map="auto", torch_dtype=torch.float16,)
+        TOKENIZER = AutoTokenizer.from_pretrained("berkeley-nest/Starling-LM-7B-alpha")
+        MODEL = AutoModelForCausalLM.from_pretrained("berkeley-nest/Starling-LM-7B-alpha")
         print("loaded model")