Update main.py

main.py

@@ -1,12 +1,12 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
+from huggingface_hub import InferenceClient
 import uvicorn
 import prompt_style
 
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 
-# client = InferenceClient(model_id)
+model_id = "failspy/Meta-Llama-3-8B-Instruct-abliterated-v3"
+client = InferenceClient(model_id)
 
 class Item(BaseModel):
     prompt: str
@@ -45,30 +45,19 @@ def generate(item: Item):
     )
 
     formatted_prompt = format_prompt(item)
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    output = ""
 
-    model_id = "failspy/Meta-Llama-3-8B-Instruct-abliterated-v3"
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto",)
-
-    input_ids = tokenizer.apply_chat_template(formatted_prompt, add_generation_prompt=True, return_tensors="pt").to(model.device)
-    
-    terminators = [
-        tokenizer.eos_token_id,
-        tokenizer.convert_tokens_to_ids("<|eot_id|>")
-    ]
-
-    outputs = model.generate(input_ids, eos_token_id=terminators, do_sample=True,  **generate_kwargs,)
-    response = outputs[0][input_ids.shape[-1]:]
-    return tokenizer.decode(response, skip_special_tokens=True)
-    
-    # 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
+    for response in stream:
+        output += response.token.text
+    return output
 
 @app.post("/generate/")
 async def generate_text(item: Item):
     ans = generate(item)
-    return {"response": ans}
+    return {"response": ans}
+
+
+@app.get("/")
+def read_root():
+    return {"Hello": "World!"}