Complete API

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-# from inference import *
+from inference import *
 
 
 st.title("💬 Chatbot")
@@ -16,8 +16,8 @@ if prompt := st.chat_input():
     st.session_state.messages.append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
 
-    response = "" #client.chat.completions.create(model="gpt-3.5-turbo", messages=st.session_state.messages)
-    msg = "" #response.choices[0].message.content
+    response = generate_text(st.session_state.messages) #client.chat.completions.create(model="gpt-3.5-turbo", messages=st.session_state.messages)
+    msg = response #response.choices[0].message.content
 
     st.session_state.messages.append({"role": "assistant", "content": msg})
     st.chat_message("assistant").write(msg)

inference.py

@@ -0,0 +1,45 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+import torch
+
+model_name = "mistralai/Mistral-7B-Instruct-v0.2"
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+)
+
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+    device_map="auto",
+    low_cpu_mem_usage=True,
+    # load_in_4bit = True,
+    quantization_config = bnb_config
+)
+
+
+
+def generate_text(messages):
+
+    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
+    no_token_encodeds = tokenizer.apply_chat_template(messages, tokenize=False).replace('<s>', "").replace('</s>', "")
+
+    output = model.generate(
+        encodeds,
+        max_length=200,
+        do_sample=True,
+        top_k=10,
+        num_return_sequences=1,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+    output_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    return output_text[len(no_token_encodeds) + 2:]
+
+    # # Remove Prompt Echo from Generated Text
+    # cleaned_output_text = output_text.replace(input_text, "")
+    # return cleaned_output_text