add claude model

README.md

@@ -10,3 +10,5 @@ pinned: false
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+ANTHROPIC_API_KEY

app.py

@@ -10,15 +10,14 @@ import shutil
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from sklearn.neighbors import NearestNeighbors
-import huggingface_hub
+import anthropic
 
-# openai.base_url = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1/v1/"
-# openai.api_key = huggingface_hub.get_token()
+# client = OpenAI(
+#     base_url='https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1/v1/',
+#     api_key=os.getenv('openai_key')
+# )
 
-clinet = OpenAI(
-    base_url='https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1/v1/',
-    api_key=os.getenv('key')
-)
+client = anthropic.Anthropic()
 
 from util import pdf_to_text, text_to_chunks, SemanticSearch
 
@@ -31,18 +30,29 @@ def load_recommender(path, start_page=1):
     return 'Corpus Loaded.'
 
 
-def generate_text(prompt, model = "gpt-3.5-turbo-16k-0613"):
-    
-    model="mistralai/Mixtral-8x7B-Instruct-v0.1"
-    max_tokens=1024
-    message = clinet.chat.completions.create(
+# def openai_generate_text(prompt, model = "gpt-3.5-turbo-16k-0613"):
+#     model="mistralai/Mixtral-8x7B-Instruct-v0.1"
+#     max_tokens=1024
+#     message = clinet.chat.completions.create(
+#         model=model,
+#         messages=[
+#             {"role": "user", "content": prompt}
+#         ],
+#         max_tokens=max_tokens,
+#     ).choices[0].message.content
+#     return message
+
+def claude_generate_text(prompt, model = "claude-3-haiku-20240307"):
+    message = client.messages.create(
         model=model,
+        max_tokens=1000,
+        temperature=0.0,
+        # system="Respond only in mandarin",
         messages=[
             {"role": "user", "content": prompt}
-        ],
-        max_tokens=max_tokens,
-    ).choices[0].message.content
-    return message
+        ]
+    )
+    return message.content[0].text
 
 def generate_answer(question):
     topn_chunks = recommender(question)
@@ -57,10 +67,10 @@ def generate_answer(question):
               "in English \n\nQuery: "
     
     prompt += f"{question}\nAnswer:"
-    answer = generate_text(prompt)
+    answer = claude_generate_text(prompt)
     return answer
 
-def question_anwser(chat_history, file, question):
+def question_answer(chat_history, file, question):
     suffix = Path(file.name).suffix
     with NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
         shutil.copyfile(file.name, tmp.name)
@@ -91,7 +101,7 @@ with gr.Blocks(css="""#chatbot { font-size: 14px; min-height: 1200; }""") as dem
             chatbot = gr.Chatbot(label="Chat History", elem_id="chatbot")
 
     btn.click(
-        question_anwser,
+        question_answer,
         inputs=[chatbot, file, question],
         outputs=[chatbot],
         api_name="predict",

requirements.txt

@@ -4,4 +4,5 @@ numpy
 scikit-learn
 tensorflow-cpu
 tensorflow-hub
-gradio
+gradio
+anthropic