Update main.py

main.py

@@ -7,7 +7,19 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 app = FastAPI()
 
-client = InferenceClient("HFHAB/FinetunedMistralModel")
+model_id = "mistralai/Mistral-7B-v0.1
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+#prompt = "<s>[INST] Write a tweet on future of AI [/INST]"
+inputs = tokenizer(prompt, return_tensors="pt").to(0)
+
+out = model.generate(**inputs, max_new_tokens=250, temperature = 0.6, top_p=0.95, tok_k=40)
+
+print(tokenizer.decode(out[0], skip_special_tokens=True))
+
+
+#client = InferenceClient("HFHAB/FinetunedMistralModel")
 
 class Item(BaseModel):
     prompt: str
@@ -46,11 +58,14 @@ def generate(item: Item):
     )
 
     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 = ""
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(0)
+    out = model.generate(**inputs, max_new_tokens=250, temperature = 0.6, top_p=0.95, tok_k=40)
+    output = tokenizer.decode(out[0], 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
+    #for response in stream:
+    #    output += response.token.text
     return output
 
 @app.post("/generate/")