Update app.py

app.py

@@ -1,4 +1,4 @@
-#from openai import OpenAI
+from openai import OpenAI
 import streamlit as st
 
 # import torch
@@ -22,34 +22,41 @@ import streamlit as st
 # hf_model = HuggingFacePipeline(pipeline=pipe
 
 
-from langchain_community.llms import HuggingFaceHub
+# from langchain_community.llms import HuggingFaceHub
 
-llm = HuggingFaceHub(
-    repo_id="google/gemma-2b-it",
-    task="text-generation",
-    model_kwargs={
-        "max_new_tokens": 512,
-        "top_k": 30,
-        "temperature": 0.1,
-        "repetition_penalty": 1.03
-    },
-)
 
+# llm = HuggingFaceHub(
+#     repo_id="google/gemma-2b-it",
+#     task="text-generation",
+#     model_kwargs={
+#         "max_new_tokens": 512,
+#         "top_k": 30,
+#         "temperature": 0.1,
+#         "repetition_penalty": 1.03
+#     },
+# )
+
+# initialize the client
+client = OpenAI(
+  base_url="https://api-inference.huggingface.co/v1",
+  #api_key=os.environ.get('HUGGINGFACEHUB_API_TOKEN')#"hf_xxx" # Replace with your token
+) 
 
-from langchain.schema import (
-    HumanMessage,
-    SystemMessage,
-)
-from langchain_community.chat_models.huggingface import ChatHuggingFace
 
-messages = [
-    SystemMessage(content="You're a helpful assistant"),
-    HumanMessage(
-        content=""
-    ),
-]
+# from langchain.schema import (
+#     HumanMessage,
+#     SystemMessage,
+# )
+# from langchain_community.chat_models.huggingface import ChatHuggingFace
 
-chat_model = ChatHuggingFace(llm=llm)
+# messages = [
+#     SystemMessage(content="You're a helpful assistant"),
+#     HumanMessage(
+#         content=""
+#     ),
+# ]
+
+#chat_model = ChatHuggingFace(llm=llm)
 
 
 # from dotenv import load_dotenv
@@ -86,9 +93,8 @@ if prompt := st.chat_input():
     #client = OpenAI()
     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)
-    res = chat_model.invoke(messages)
-    #msg = response.choices[0].message.content
-    msg = res.content
+    response = client.chat.completions.create(model="google/gemma-2b-it", messages=st.session_state.messages)
+
+    msg = response.choices[0].message.content
     st.session_state.messages.append({"role": "assistant", "content": msg})
     st.chat_message("assistant").write(msg)