Update app.py

app.py

@@ -1,37 +1,34 @@
-import gradio as gr
-from transformers import AutoTokenizer, FastLanguageModel
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 
 # Load the model and tokenizer
-model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name="lora_model",  # Replace with your trained model name
-    max_seq_length=512,
-    dtype="float16",
-    load_in_4bit=True,
-)
-FastLanguageModel.for_inference(model)
+model_name = "Rafay17/Llama3.2_1b_customModel2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")  # Ensure to load the model on GPU
 
-# Define the inference function
-def generate_response(user_input):
-    # Prepare the input for the model
-    labeled_prompt = (
-        "Please provide the response with the following labels:\n"
-        f"User Input: {user_input}\n"
-        "Response:"
-    )
+# Prepare the model for inference
+model.eval()
 
-    inputs = tokenizer(
-        [labeled_prompt],
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=512,
-    ).to("cuda")
+# Define a function to generate responses
+def generate_response(input_text):
+    # Prepare the input for the model
+    inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
 
-    response = model.generate(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
-    return tokenizer.decode(response[0], skip_special_tokens=True)
+    # Set up the text streamer to stream the generated response
+    text_streamer = TextStreamer(tokenizer, skip_prompt=True)
 
-# Create a Gradio interface
-iface = gr.Interface(fn=generate_response, inputs="text", outputs="text", title="Chatbot Interface", description="Enter your message below:")
+    # Generate the response
+    with torch.no_grad():
+        model.generate(
+            input_ids=inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            streamer=text_streamer,
+            max_new_tokens=64,  # Adjust this value as needed
+            pad_token_id=tokenizer.eos_token_id,
+        )
 
-# Launch the app
-iface.launch()
+# Example usage of the generate_response function
+input_text = "Hello, how can I help you today?"
+print("Generating response for input:")
+print(input_text)
+generate_response(input_text)