Update app.py

app.py

@@ -4,43 +4,29 @@ from peft import PeftModel  # Ensure PEFT is installed: pip install peft
 import os
 
 # Define the model and base paths
-model_path = "FadQ/gemma-2b-diary-consultaton-chatbot"
-base_model = "google/gemma-2b"
+model = "FadQ/gemma-2b-diary-consultaton-chatbot"
+# base_model = "google/gemma-2b"
 
-# Use your Hugging Face token
-hf_token = os.getenv('HF_TOKEN')
+# # Use your Hugging Face token
+# hf_token = os.getenv('HF_TOKEN')
 
-# Load tokenizer with authentication
-tokenizer = AutoTokenizer.from_pretrained(base_model, token=hf_token, force_download=True)
+# # Load tokenizer with authentication
+# tokenizer = AutoTokenizer.from_pretrained(base_model, token=hf_token, force_download=True)
 
-# Load the base model and apply adapter with authentication
-base_model = AutoModelForCausalLM.from_pretrained(base_model, device_map="auto", token=hf_token)
-model = PeftModel.from_pretrained(base_model, model_path)
+# # Load the base model and apply adapter with authentication
+# base_model = AutoModelForCausalLM.from_pretrained(base_model, device_map="auto", token=hf_token)
+# model = PeftModel.from_pretrained(base_model, model_path)
 
 
 # Create pipeline
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=0  # Assuming GPU is available
-)
+pipe = pipeline("text-generation", model=model_path, device=0)
 
-def predict(input_text, system_message, max_new_tokens, temperature, top_p):
-    # Format the prompt
-    prompt = f"{system_message}\nUser: {input_text}\nAssistant:"
-
-    # Generate text using the pipeline
-    result = pipe(
-        prompt,
-        max_length=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        num_return_sequences=1
-    )
+def predict(input_text):
+    result = pipe(input_text, max_length=150, num_return_sequences=1)
     generated_text = result[0]["generated_text"]
     return generated_text
 
+
 # Create the Gradio interface
 demo = gr.Interface(
     fn=predict,