Update src/model_consumer.py

src/model_consumer.py

@@ -7,7 +7,8 @@ model_id = "prd101-wd/phi1_5-bankingqa-merged"
 # Load model only once
 @st.cache_resource
 def load_model():
-    return pipeline("question-answering", model=model_id)
+    #return pipeline("question-answering", model=model_id)
+    return pipeline("text-generation", model=model_id, trust_remote_code=True)
                     
 # Create a text generation pipeline
 pipe = load_model()
@@ -23,10 +24,31 @@ if st.button("Ask"):
     if user_input.strip():
         # Format the prompt like Alpaca-style
         prompt = f"### Instruction:\n{user_input}\n\n### Response:\n"
-        output = pipe(prompt, max_new_tokens=200, do_sample=True)[0]["generated_text"]
+        output = pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
+
+         # Process output
+        if isinstance(output, list) and output:
+            answer = output[0]['generated_text']
+            # Extract only the response part
+            if "### Response:" in answer:
+                answer = answer.split("### Response:")[-1].strip()
+        else:
+            answer = "Unable to generate a response. Please try again."
+            
+
+        
+        # if isinstance(output, list) and len(output) > 0 and "generated_text" in output[0]:
+        #     answer = output[0]["generated_text"]
+        # else:
+        #     answer = "Unable to generate a response. Please try again."
 
         # Extract only the model's response (remove prompt part if included in output)
-        answer = output.split("### Response:")[-1].strip()
+        #answer = output.split("### Response:")[-1].strip()
+        # if isinstance(output, str):
+        #     answer = output.split("### Response:")[-1].strip()
+        # else:
+        #     answer = "Unexpected output format. Please try again."
+        
         st.markdown("### HelpdeskBot Answer:")
         st.success(answer)
     else: