Update app.py

app.py

@@ -1,57 +1,56 @@
-from transformers import AutoTokenizer, TextStreamer
+import gradio as gr
 from unsloth import FastLanguageModel
-import torch
+from transformers import AutoTokenizer, TextStreamer
 
 # Load the model and tokenizer
-model_name = "Rafay17/Llama3.2_1b_customModle2"  # Your model name
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-model = FastLanguageModel.from_pretrained(
-    model_name=model_name,
-    max_seq_length=512,  # Adjust as needed
-    dtype="float16",     # Adjust as needed
-    load_in_4bit=True    # Adjust based on your needs
-)
-
-FastLanguageModel.for_inference(model)  # Call this immediately after loading the model
+model_name = "Rafay17/Llama3.2_1b_customModel2"  # Your custom model
+model, tokenizer = FastLanguageModel.from_pretrained(model_name)
+FastLanguageModel.for_inference(model)  # Enable the model for inference
 
 # Function to generate a response
-def generate_response(input_text):
-    # Prepare the labeled prompt for the model
-    labeled_prompt = f"User Input: {input_text}\nResponse:"
+def generate_response(message, history, max_tokens, temperature, top_p):
+    # Prepare the labeled prompt for response generation
+    labeled_prompt = f"User Input: {message}\nResponse:"
 
-    # Prepare the input for the model
+    # Tokenize the input
     inputs = tokenizer(
         [labeled_prompt],
         return_tensors="pt",
         padding=True,
         truncation=True,
-        max_length=512,  # Ensure this matches your model's max length
+        max_length=512,
     ).to("cuda")
 
-    # Set up the text streamer to stream the generated response
+    # Generate the response
     text_streamer = TextStreamer(tokenizer, skip_prompt=True)
+    response = ""
+    for token in model.generate(
+        input_ids=inputs.input_ids,
+        attention_mask=inputs.attention_mask,
+        streamer=text_streamer,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        pad_token_id=tokenizer.eos_token_id,
+    ):
+        response += token
+
+    return response
+
+
+# Define the Gradio interface
+demo = gr.Interface(
+    fn=generate_response,
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Enter your message here..."),
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Slider(minimum=1, maximum=512, value=64, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, label="Top-p (nucleus sampling)"),
+    ],
+    outputs=gr.Textbox(label="Chatbot Response"),
+    live=True
+)
 
-    # Generate the response
-    with torch.no_grad():  # Disable gradient calculation for inference
-        model.generate(
-            input_ids=inputs.input_ids,
-            attention_mask=inputs.attention_mask,
-            streamer=text_streamer,
-            max_new_tokens=100,  # Adjust this value as needed
-            pad_token_id=tokenizer.eos_token_id,
-        )
-
-# Function to take user input and generate output
-def user_interaction():
-    print("Welcome to the Chatbot! Type 'exit' to quit.")
-    while True:
-        user_input = input("You: ")
-        if user_input.lower() == 'exit':
-            print("Exiting the chatbot. Goodbye!")
-            break
-        print("Chatbot is generating a response...")
-        generate_response(user_input)
-
-# Start the user interaction
-user_interaction()
+if __name__ == "__main__":
+    demo.launch()