Update main.py

main.py

@@ -3,20 +3,57 @@ from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import uvicorn
 
+# **************************************************
+# import transformers
+# import torch
+
+model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+
+# pipeline = transformers.pipeline(
+#     "text-generation",
+#     model=model_id,
+#     model_kwargs={"torch_dtype": torch.bfloat16},
+#     device_map="auto",
+# )
+
+def generate(item: Item):
+    messages = [
+        {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
+        {"role": "user", "content": "Who are you?"},
+    ]
+    
+    terminators = [
+        pipeline.tokenizer.eos_token_id,
+        pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
+    ]
+    
+    outputs = pipeline(
+        messages,
+        max_new_tokens=item.max_new_tokens,
+        eos_token_id=terminators,
+        do_sample=True,
+        temperature=item.temperature,
+        top_p=item.top_p,
+    )
+    return outputs[0]["generated_text"][-1]
+
+# **************************************************
 
-app = FastAPI()
 
-client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+
+
+client = InferenceClient(model_id)
 
 class Item(BaseModel):
     prompt: str
     history: list
     system_prompt: str
-    temperature: float = 0.7
-    max_new_tokens: int = 512
-    top_p: float = 0.15
-    repetition_penalty: float = 1.0
-    seed: int = 42
+    temperature: float = 0.6
+    max_new_tokens: int = 1024
+    top_p: float = 0.95
+    seed : int = 42
+    
+app = FastAPI()
 
 def format_prompt(message, history):
     prompt = "<s>"