llama3

app.py

@@ -1,15 +1,22 @@
 from fastapi import FastAPI
+from fastapi.responses import RedirectResponse
 from transformers import pipeline
- 
+import torch
+from huggingface_hub import hf_hub_download
+
+hf_hub_download(repo_id="google/pegasus-xsum", filename="config.json")
+
 # Create a new FastAPI app instance
 app = FastAPI()
  
+model_id = "meta-llama/Meta-Llama-3-70B"
+
 # Initialize the text generation pipeline
 # This function will be able to generate text
 # given an input.
-pipe = pipeline("text2text-generation", 
-model="google/flan-t5-small")
- 
+pipe = pipeline("text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto")
+#pipe = pipeline("text2text-generation", model="google/flan-t5-small")
+
 # Define a function to handle the GET request at `/generate`
 # The generate() function is defined as a FastAPI route that takes a 
 # string parameter called text. The function generates text based on the # input using the pipeline() object, and returns a JSON response 
@@ -25,4 +32,8 @@ def generate(text: str):
     output = pipe(text)
      
     # Return the generated text in a JSON response
-    return {"output": output[0]["generated_text"]}
+    return {"output": output[0]["generated_text"]}
+
+@app.get("/")
+async def docs_redirect():
+    return RedirectResponse(url='/docs')