Update app.py

app.py

@@ -1,44 +1,70 @@
-from transformers import pipeline
-from gradio import gr
-from huggingface_hub import InferenceClient
+import logging
+import os
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-pipe = pipeline("fill-mask", model="google-bert/bert-base-uncased")
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-def respond(user_input, history, system_message, max_tokens, temperature, top_p):
-    messages = [{"role": "system", "content": system_message}]
+# Checkpoint paths
+model_checkpoint_path = "model_checkpoint.pth"
+tokenizer_checkpoint_path = "tokenizer_checkpoint.pth"
 
-    for val in history:
-        if val["role"] == "user":
-            messages.append(val)
-        if val["role"] == "assistant":
-            messages.append(val)
+# Load model and tokenizer from checkpoint if they exist
+if os.path.exists(model_checkpoint_path) and os.path.exists(tokenizer_checkpoint_path):
+    try:
+        model = torch.load(model_checkpoint_path)
+        tokenizer = torch.load(tokenizer_checkpoint_path)
+        logger.info("Model and tokenizer loaded from checkpoint.")
+    except Exception as e:
+        logger.error(f"Failed to load model or tokenizer from checkpoint: {e}")
+        raise
+else:
+    # Load model directly
+    try:
+        tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
+        model = AutoModelForCausalLM.from_pretrained("google/gemma-2-9b")
+        logger.info("Model and tokenizer loaded successfully.")
+    except Exception as e:
+        logger.error(f"Failed to load model or tokenizer: {e}")
+        raise
 
+def respond(user_input, history, system_message, max_tokens=20, temperature=0.9, top_p=0.9):
+    messages = [{"role": "system", "content": system_message}]
+    messages.extend(history)
     messages.append({"role": "user", "content": user_input})
 
-    response = ""
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
+    # Convert messages to a single string
+    input_text = " ".join([msg["content"] for msg in messages])
+
+    # Tokenize the input text
+    inputs = tokenizer(input_text, return_tensors="pt")
+
+    # Generate attention mask
+    attention_mask = inputs["attention_mask"]
+
+    # Generate text using the model
+    outputs = model.generate(
+        inputs.input_ids,
+        attention_mask=attention_mask,
+        max_length=max_tokens,
         temperature=temperature,
         top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        yield response
-
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-demo = gr.ChatInterface(
-    respond,
-    inputs=[
-        gr.Textbox(placeholder="Type your message here", label="User input"),
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-    ],
-)
+        pad_token_id=tokenizer.eos_token_id,
+        do_sample=True
+    )
+
+    # Decode the generated text
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return response
 
 if __name__ == "__main__":
-    demo.launch()
+    print("Welcome to the Chatbot!")
+    while True:
+        user_input = input("You: ")
+        system_message = "Chatbot: "
+        history = [{"role": "assistant", "content": "Hello, how can I assist you today?"}]
+        response = respond(user_input, history, system_message)
+        print(response)