Update app.py

app.py

@@ -1,10 +1,21 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, pipeline
+import torch
+
+# Define the model repository
+REPO_NAME = 'schuler/experimental-JP47D20'
+# REPO_NAME = 'schuler/experimental-JP47D21-KPhi-3-micro-4k-instruct'
+
+# How to cache?
+def load_model(repo_name):
+    tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
+    generator_conf = GenerationConfig.from_pretrained(repo_name)
+    model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True, torch_dtype=torch.bfloat16)
+    return tokenizer, generator_conf, model
+
+tokenizer, generator_conf, model = load_model(REPO_NAME)
 
 
 def respond(
@@ -25,8 +36,29 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    response = ""
+    for message in messages:
+            role = "<|assistant|>" if message['role'] == 'assistant' else "<|user|>"
+            prompt += f"\n{role}\n{message['content']}\n<|end|>\n"        
+    # prompt += f"\n<|user|>\n{user_text}\n<|end|><|assistant|>\n"
+
+    # Generate the response
+    response_output = generator(
+        prompt,
+        generation_config=generator_conf,
+        max_new_tokens=64,
+        do_sample=True,
+        top_p=0.25,
+        repetition_penalty=1.2
+    )
+
+    generated_text = response_output[0]['generated_text']
+
+    # st.session_state.last_response = generated_text
 
+    # Extract the assistant's response
+    yield generated_text[len(prompt):].strip()
+    
+    """
     for message in client.chat_completion(
         messages,
         max_tokens=max_tokens,
@@ -38,6 +70,7 @@ def respond(
 
         response += token
         yield response
+    """
 
 
 """