Update app.py

app.py

@@ -2,7 +2,6 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 
 app = FastAPI()
 app.add_middleware(
@@ -13,9 +12,13 @@ app.add_middleware(
     allow_headers=["*"]
 )
 
-model_name = "togethercomputer/RedPajama-INCITE-7B-Base"
+model_name = "ethzanalytics/RedPajama-INCITE-7B-Base-sharded-bf16"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="bfloat16",
+    device_map="auto"
+)
 
 class PromptRequest(BaseModel):
     prompt: str
@@ -26,13 +29,13 @@ async def generate_story(req: PromptRequest):
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt must not be empty")
 
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
     outputs = model.generate(
         **inputs,
         max_new_tokens=200,
         do_sample=True,
+        temperature=0.9,
         top_p=0.9,
-        temperature=0.85,
         repetition_penalty=1.2
     )
     story = tokenizer.decode(outputs[0], skip_special_tokens=True)