app.py

import streamlit as st
import torch

from transformers import AutoTokenizer, AutoModelForCausalLM


device = torch.device("cpu")

# Load the model and tokenizer from Hugging Face
model_name = "anjikum/ph2-sft-retrained"  # Replace with your model's Hugging Face repo name
model = AutoModelForCausalLM.from_pretrained(model_name,load_in_8bit=False).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Streamlit UI elements
st.title("Phi-2 Fine-Tuned Model")
st.write("Input a prompt and the model will generate a response.")

# User input
prompt = st.text_area("Enter your prompt:")

if st.button("Generate Answer"):
    if prompt:
        # Tokenize the prompt and generate a response
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        outputs = model.generate(**inputs, max_length=100, num_return_sequences=1)
        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Display the answer
        st.write("Generated Answer:")
        st.write(answer)
    else:
        st.warning("Please enter a prompt!")