Update main.py

main.py

@@ -1,7 +1,7 @@
 #pip install fastapi
 #uvicorn main:app --reload
 #import gradio as gr
-
+import torch
 from transformers import pipeline
 from fastapi import FastAPI
 
@@ -13,6 +13,22 @@ generator = pipeline("text-generation", model="TheBloke/zephyr-7B-alpha-GGUF")
 #model = AutoModel.from_pretrained("TheBloke/zephyr-7B-alpha-GGUF")
 
 
+
+pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-alpha", torch_dtype=torch.bfloat16, device_map="auto")
+
+# We use the tokenizer's chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
+messages = [
+    {
+        "role": "system",
+        "content": "You are a Spiritual Coach who always responds in the most profound and poetic style",
+    },
+    {"role": "user", "content": "What is Life?"},
+]
+prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+outputs = pipe(prompt, max_new_tokens=2560, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+print(outputs[0]["generated_text"])
+
+
 @app.get("/")
 async def root():
     return {"message": "Hello World"}
@@ -21,4 +37,5 @@ async def root():
 @app.post("/predict")
 async def root(text):
     #return {"message": "Hello World"}
-    return generator(text,max_length=100, num_return_sequences=1)
+    #return generator(text,max_length=2560, num_return_sequences=1)
+    return outputs[0]["generated_text"]