solar

question_paper.py

@@ -12,15 +12,12 @@ TOKENIZER = None
 
 @app.get("/")
 def llama():
-    text = "Hi, my name is "
-    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)
+    prompt = [{'role': 'user', 'content': 'List 3 synonyms for the word "tiny"'}]
+    inputs = TOKENIZER.apply_chat_template( prompt, add_generation_prompt=True,     return_tensors='pt' )
+
+    tokens = MODEL.generate( inputs.to(MODEL.device), max_new_tokens=1024, temperature=0.3, do_sample=True)
+    
+    tresponse = TOKENIZER.decode(tokens[0], skip_special_tokens=False)
     print(tresponse)
 
     return tresponse
@@ -31,6 +28,6 @@ def init_model():
     global TOKENIZER
     if not MODEL:
         print("loading model")
-        TOKENIZER = AutoTokenizer.from_pretrained("berkeley-nest/Starling-LM-7B-alpha")
-        MODEL = AutoModelForCausalLM.from_pretrained("berkeley-nest/Starling-LM-7B-alpha")
+        TOKENIZER = AutoTokenizer.from_pretrained('stabilityai/stablelm-zephyr-3b')
+        MODEL = AutoModelForCausalLM.from_pretrained('stabilityai/stablelm-zephyr-3b', device_map="auto")
         print("loaded model")