Update app.py

app.py

@@ -18,23 +18,23 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 #     pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
 #     result = pipe(f"<s>[INST] {input_text} [/INST]")
 #     return result[0]['generated_text']
-
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from fastapi import FastAPI
 
-app = FastAPI()
+# Specify the path to your fine-tuned model and tokenizer
+model_path = "./"  # Assuming the model is in the same directory as your notebook
+model_name = "pytorch_model-00001-of-00002.bin"  # Replace with your model name
 
-model_name = "pytorch_model-00001-of-00002.bin"  # Replace with your Hugging Face model name
+# Load the model and tokenizer
+model = AutoModelForCausalLM.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
 
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Example usage
+input_text = "Once upon a time"
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+output = model.generate(input_ids, max_length=50, num_return_sequences=1)
+generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
-@app.post("/generate/")
-async def generate_text(prompt: str):
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    output = model.generate(input_ids, max_length=50, num_return_sequences=1)
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"generated_text": generated_text}
+print(generated_text)
 
 
 text_generation_interface = gr.Interface(