Update main.py

main.py

@@ -7,59 +7,59 @@
 #         "history": [],
 #         "system_prompt": "You are a very powerful AI assistant."
 #       }' \
-#   https://iiced-mixtral-46-7b-fastapi.hf.space/generate/
-
+#   https://phk0-bai.hf.space/generate/
 
 from fastapi import FastAPI
 from pydantic import BaseModel
-from huggingface_hub import InferenceClient
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import uvicorn
+import torch
 
 
 app = FastAPI()
 
-client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
-
 class Item(BaseModel):
     prompt: str
     history: list
     system_prompt: str
     temperature: float = 0.0
-    max_new_tokens: int = 1048
+    max_new_tokens: int = 900
     top_p: float = 0.15
     repetition_penalty: float = 1.0
 
-def format_prompt(message, history):
-    prompt = "<s>"
+def format_prompt(system, message, history):
+    prompt = [{"role": "system", "content": system}] 
     for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
+        prompt += {"role": "user", "content": user_prompt}
+        prompt += {"role": "assistant", "content": bot_response}
+    prompt += {"role": "user", "content": message}
     return prompt
 
 def generate(item: Item):
-    temperature = float(item.temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(item.top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=item.max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=item.repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model_path = "ibm-granite/granite-34b-code-instruct-8k"
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    # drop device_map if running on CPU
+    model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
+    model.eval()
+    # change input text as desired
+    chat = format_prompt(item.system_prompt, item.prompt, item.history)
+    chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+    # tokenize the text
+    input_tokens = tokenizer(chat, return_tensors="pt")
+    # transfer tokenized inputs to the device
+    for i in input_tokens:
+        input_tokens[i] = input_tokens[i].to(device)
+    # generate output tokens
+    output = model.generate(**input_tokens, max_new_tokens=900)
+    output_text = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+    return output_text
 
-    formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-    return output
 
 @app.post("/generate/")
 async def generate_text(item: Item):
     return {"response": generate(item)}
+
+@app.get("/")
+async def generate_text_root(item: Item):
+    return {"response": "try entry point: /generate/"}